69967b25c6
OllamaClient.Generate/GenerateText/call now take ctx and build the HTTP request via NewRequestWithContext, so an in-flight LLM call is aborted when the parent context is cancelled (Ctrl-C). Classifier.Classify and its tier helpers take ctx too. ProcessFunc gets a ctx parameter so the poller can forward its cancellable context down to classification. Explainer.summarize manages its own 60s context since reaction-driven flow has no parent ctx to inherit.
121 lines
3.2 KiB
Go
121 lines
3.2 KiB
Go
package classifier
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"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. Honors ctx cancellation — callers can abort an in-flight
|
|
// call instead of waiting for the configured timeout.
|
|
func (o *OllamaClient) Generate(ctx context.Context, prompt string) (string, error) {
|
|
return o.call(ctx,
|
|
"You are a JSON-only news classification API. Respond with valid JSON only. No preamble, no explanation, no markdown.",
|
|
prompt,
|
|
"json",
|
|
)
|
|
}
|
|
|
|
// GenerateText sends a system + user prompt to Ollama and returns the raw
|
|
// text response. Unlike Generate, no JSON mode is forced.
|
|
func (o *OllamaClient) GenerateText(ctx context.Context, systemPrompt, userPrompt string) (string, error) {
|
|
return o.call(ctx, systemPrompt, userPrompt, "")
|
|
}
|
|
|
|
func (o *OllamaClient) call(ctx context.Context, systemPrompt, userPrompt, format string) (string, error) {
|
|
reqBody := chatRequest{
|
|
Model: o.model,
|
|
Messages: []chatMessage{
|
|
{Role: "system", Content: systemPrompt},
|
|
{Role: "user", Content: userPrompt},
|
|
},
|
|
Stream: false,
|
|
Format: format,
|
|
}
|
|
|
|
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)
|
|
|
|
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
|
|
if err != nil {
|
|
return "", fmt.Errorf("ollama request build: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
resp, err := o.httpClient.Do(req)
|
|
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
|
|
}
|