package classifier import ( "bytes" "encoding/json" "fmt" "io" "log/slog" "net/http" "strings" "time" "pete/internal/config" ) // OllamaClient handles communication with the Ollama API. type OllamaClient struct { baseURL string model string httpClient *http.Client } type chatRequest struct { Model string `json:"model"` Messages []chatMessage `json:"messages"` Stream bool `json:"stream"` Format string `json:"format"` } type chatMessage struct { Role string `json:"role"` Content string `json:"content"` } type chatResponse struct { Message struct { Content string `json:"content"` } `json:"message"` } // NewOllamaClient creates an Ollama API client from config. func NewOllamaClient(cfg config.OllamaConfig) *OllamaClient { return &OllamaClient{ baseURL: strings.TrimRight(cfg.BaseURL, "/"), model: cfg.Model, httpClient: &http.Client{ Timeout: time.Duration(cfg.TimeoutSeconds) * time.Second, }, } } // Generate sends a prompt to Ollama via the chat API and returns the raw response text. func (o *OllamaClient) Generate(prompt string) (string, error) { reqBody := chatRequest{ Model: o.model, Messages: []chatMessage{ { Role: "system", Content: "You are a JSON-only news classification API. Respond with valid JSON only. No preamble, no explanation, no markdown.", }, { Role: "user", Content: prompt, }, }, Stream: false, Format: "json", } body, err := json.Marshal(reqBody) if err != nil { return "", fmt.Errorf("marshal request: %w", err) } url := o.baseURL + "/api/chat" slog.Debug("ollama: calling", "url", url, "model", o.model) resp, err := o.httpClient.Post(url, "application/json", bytes.NewReader(body)) if err != nil { return "", fmt.Errorf("ollama request: %w", err) } defer resp.Body.Close() respBody, err := io.ReadAll(io.LimitReader(resp.Body, 1*1024*1024)) // 1MB cap if err != nil { return "", fmt.Errorf("read ollama response: %w", err) } if resp.StatusCode != http.StatusOK { return "", fmt.Errorf("ollama status %d: %s", resp.StatusCode, string(respBody)) } var chatResp chatResponse if err := json.Unmarshal(respBody, &chatResp); err != nil { return "", fmt.Errorf("decode ollama response: %w (body: %s)", err, truncate(string(respBody), 200)) } result := strings.TrimSpace(chatResp.Message.Content) if result == "" { return "", fmt.Errorf("ollama returned empty response (body: %s)", truncate(string(respBody), 200)) } return result, nil }