play-life/play-life-llm/internal/ollama/client.go

package ollama

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

const defaultTimeout = 10 * time.Minute

// Client calls Ollama /api/chat.
type Client struct {
	BaseURL    string
	HTTPClient *http.Client
}

// NewClient creates an Ollama client. baseURL is e.g. "http://localhost:11434".
func NewClient(baseURL string) *Client {
	return &Client{
		BaseURL: baseURL,
		HTTPClient: &http.Client{
			Timeout: defaultTimeout,
		},
	}
}

// ChatRequest matches Ollama POST /api/chat body.
type ChatRequest struct {
	Model    string        `json:"model"`
	Messages []ChatMessage `json:"messages"`
	Stream   bool          `json:"stream"`
	Format   interface{}   `json:"format,omitempty"` // "json" or JSON schema object
	Tools    []Tool        `json:"tools,omitempty"`
}

// ChatMessage is one message in the conversation.
type ChatMessage struct {
	Role       string        `json:"role"` // "user", "assistant", "system", "tool"
	Content    string        `json:"content,omitempty"`
	ToolCalls  []ToolCall    `json:"tool_calls,omitempty"`
	ToolName   string        `json:"tool_name,omitempty"` // for role "tool"
}

// Tool defines a function the model may call.
type Tool struct {
	Type     string   `json:"type"`
	Function ToolFunc `json:"function"`
}

// ToolFunc describes the function.
type ToolFunc struct {
	Name        string      `json:"name"`
	Description string      `json:"description"`
	Parameters  interface{}  `json:"parameters"`
}

// ToolCall is a model request to call a tool.
type ToolCall struct {
	Type     string     `json:"type"`
	Function ToolCallFn `json:"function"`
}

// ToolCallFn holds name and arguments.
// Arguments may come from Ollama as a JSON object or as a JSON string.
type ToolCallFn struct {
	Name      string      `json:"name"`
	Arguments interface{} `json:"arguments"` // object or string
}

// QueryFromToolCall returns the "query" argument from a web_search tool call.
// Ollama may send arguments as a map or as a JSON string.
func QueryFromToolCall(tc ToolCall) string {
	switch v := tc.Function.Arguments.(type) {
	case map[string]interface{}:
		if q, _ := v["query"].(string); q != "" {
			return q
		}
	case string:
		var m map[string]interface{}
		if json.Unmarshal([]byte(v), &m) == nil {
			if q, _ := m["query"].(string); q != "" {
				return q
			}
		}
	}
	return ""
}

// ChatResponse is the Ollama /api/chat response.
type ChatResponse struct {
	Message ChatMessage `json:"message"`
	Done    bool        `json:"done"`
}

// Chat sends a chat request and returns the response.
func (c *Client) Chat(req *ChatRequest) (*ChatResponse, error) {
	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("marshal request: %w", err)
	}
	url := c.BaseURL + "/api/chat"
	httpReq, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("new request: %w", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := c.HTTPClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("do request: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("ollama returned %d: %s", resp.StatusCode, string(b))
	}

	var out ChatResponse
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return nil, fmt.Errorf("decode response: %w", err)
	}
	return &out, nil
}

// WebSearchTool returns the tool definition for web_search (Tavily).
func WebSearchTool() Tool {
	return Tool{
		Type: "function",
		Function: ToolFunc{
			Name:        "web_search",
			Description: "Search the web for current information. Use when you need up-to-date or factual information from the internet.",
			Parameters: map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"query": map[string]interface{}{
						"type":        "string",
						"description": "Search query",
					},
				},
				"required": []string{"query"},
			},
		},
	}
}