105 lines
2.5 KiB
Go
105 lines
2.5 KiB
Go
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
|
|
}
|