Weather Tool

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
)

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type ToolDef struct {
	Type     string `json:"type"`
	Function struct {
		Name        string          `json:"name"`
		Description string          `json:"description"`
		Parameters  json.RawMessage `json:"parameters"`
	} `json:"function"`
}

type ToolCall struct {
	Function struct {
		Name      string          `json:"name"`
		Arguments json.RawMessage `json:"arguments"`
	} `json:"function"`
}

type ChatResponse struct {
	Message struct {
		Role      string     `json:"role"`
		Content   string     `json:"content"`
		ToolCalls []ToolCall `json:"tool_calls,omitempty"`
	} `json:"message"`
}

func callOllama(messages []Message, tools []ToolDef) ChatResponse {
	body, _ := json.Marshal(map[string]any{
		"model":    "llama3.2",
		"messages": messages,
		"stream":   false,
		"tools":    tools,
	})
	resp, err := http.Post("http://localhost:11434/api/chat",
		"application/json", bytes.NewReader(body))
	if err != nil {
		fmt.Println("Error:", err)
		return ChatResponse{}
	}
	defer resp.Body.Close()
	data, _ := io.ReadAll(resp.Body)
	var result ChatResponse
	json.Unmarshal(data, &result)
	return result
}

func executeToolCall(tc ToolCall) string {
	switch tc.Function.Name {
	case "get_weather":
		var args struct {
			City string `json:"city"`
		}
		json.Unmarshal(tc.Function.Arguments, &args)
		// Hardcoded for this example
		return fmt.Sprintf(`{"city": %q, "temp": 18, "condition": "cloudy"}`, args.City)
	default:
		return fmt.Sprintf(`{"error": "unknown tool: %s"}`, tc.Function.Name)
	}
}

func main() {
	// Step 1: Define the tool
	weatherTool := ToolDef{Type: "function"}
	weatherTool.Function.Name = "get_weather"
	weatherTool.Function.Description = "Get current weather for a city"
	weatherTool.Function.Parameters = json.RawMessage(`{
		"type": "object",
		"properties": {
			"city": {"type": "string", "description": "City name"}
		},
		"required": ["city"]
	}`)
	tools := []ToolDef{weatherTool}

	// Step 2: Send user message with tools
	messages := []Message{
		{Role: "user", Content: "What's the weather in Tokyo?"},
	}
	fmt.Println("User:", messages[0].Content)

	result := callOllama(messages, tools)

	// Step 3: Check if the model wants to call a tool
	if len(result.Message.ToolCalls) == 0 {
		fmt.Println("Model:", result.Message.Content)
		return
	}

	tc := result.Message.ToolCalls[0]
	fmt.Printf("Model wants to call: %s(%s)\n", tc.Function.Name, tc.Function.Arguments)

	// Step 4: Execute the tool
	toolResult := executeToolCall(tc)
	fmt.Println("Tool result:", toolResult)

	// Step 5: Send result back to get final response
	messages = append(messages, Message{Role: "assistant", Content: ""})
	messages = append(messages, Message{Role: "tool", Content: toolResult})

	final := callOllama(messages, tools)
	fmt.Println("Model:", final.Message.Content)

	// User: What's the weather in Tokyo?
	// Model wants to call: get_weather({"city":"Tokyo"})
	// Tool result: {"city": "Tokyo", "temp": 18, "condition": "cloudy"}
	// Model: It's currently 18°C and cloudy in Tokyo.
}