Initial commit

internal/classifier/ollama.go

@@ -0,0 +1,104 @@
+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
+}