velox-bot/internal/llm/ollama.go

// Package llm is a minimal client for a local Ollama server's chat API.
package llm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"strings"
)

const (
	// numCtx must match what the chunk budget in internal/tldr assumes.
	numCtx      = 8192
	numPredict  = 550
	temperature = 0.3
)

type Client struct {
	baseURL string
	model   string
	threads int
	http    *http.Client
}

// New builds a client. threads caps how many CPU threads the model may use
// (0 lets Ollama decide). Always set it on a shared box: Ollama otherwise
// starts one thread per core and they busy-wait, which both starves other
// services and, under a container CPU quota, throttles inference to a crawl.
func New(baseURL, model string, threads int) *Client {
	return &Client{
		baseURL: strings.TrimRight(baseURL, "/"),
		model:   model,
		threads: threads,
		// No client-level timeout: CPU inference can legitimately take
		// minutes, so callers bound it with the context instead.
		http: &http.Client{},
	}
}

type chatMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type chatRequest struct {
	Model    string         `json:"model"`
	Messages []chatMessage  `json:"messages"`
	Stream   bool           `json:"stream"`
	Think    *bool          `json:"think,omitempty"`
	Options  map[string]any `json:"options"`
}

type chatResponse struct {
	Message chatMessage `json:"message"`
	Done    bool        `json:"done"`
	Error   string      `json:"error"`
}

var thinkTags = regexp.MustCompile(`(?s)<think>.*?</think>`)

// answer strips any inline reasoning (some models still emit it) so the user
// only gets the reply.
func answer(s string) string {
	return strings.TrimSpace(thinkTags.ReplaceAllString(s, ""))
}

// visible is answer for text that is still being written: an unfinished
// <think> block is hidden from where it starts.
func visible(s string) string {
	s = thinkTags.ReplaceAllString(s, "")
	if i := strings.Index(s, "<think>"); i >= 0 {
		s = s[:i]
	}
	return strings.TrimSpace(s)
}

func statusError(status int, body []byte) error {
	return fmt.Errorf("ollama returned %d: %s", status, truncate(string(body), 300))
}

// post sends one request and returns the raw response.
func (c *Client) post(ctx context.Context, system, user string, stream bool, think *bool) (*http.Response, error) {
	options := map[string]any{
		"num_ctx":     numCtx,
		"num_predict": numPredict,
		"temperature": temperature,
	}
	if c.threads > 0 {
		options["num_thread"] = c.threads
	}

	body, err := json.Marshal(chatRequest{
		Model: c.model,
		Messages: []chatMessage{
			{Role: "system", Content: system},
			{Role: "user", Content: user},
		},
		Stream:  stream,
		Think:   think,
		Options: options,
	})
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/api/chat", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := c.http.Do(req)
	if err != nil {
		return nil, fmt.Errorf("ollama unreachable: %w", err)
	}
	return resp, nil
}

// open returns a 200 response ready to read. Thinking is switched off because
// a summary has no use for a reasoning preamble and on CPU it can cost most of
// the time budget. Models that don't support the flag reject it, in which
// case the request is repeated without it.
func (c *Client) open(ctx context.Context, system, user string, stream bool) (*http.Response, error) {
	noThink := false
	resp, err := c.post(ctx, system, user, stream, &noThink)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode == http.StatusBadRequest {
		raw, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		if !strings.Contains(strings.ToLower(string(raw)), "think") {
			return nil, statusError(resp.StatusCode, raw)
		}
		if resp, err = c.post(ctx, system, user, stream, nil); err != nil {
			return nil, err
		}
	}

	if resp.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		return nil, statusError(resp.StatusCode, raw)
	}
	return resp, nil
}

// Chat sends one system+user exchange and returns the model's whole reply.
func (c *Client) Chat(ctx context.Context, system, user string) (string, error) {
	resp, err := c.open(ctx, system, user, false)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	var parsed chatResponse
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return "", fmt.Errorf("decoding ollama response: %w", err)
	}
	if parsed.Error != "" {
		return "", fmt.Errorf("ollama error: %s", parsed.Error)
	}

	out := answer(parsed.Message.Content)
	if out == "" {
		return "", fmt.Errorf("ollama returned an empty reply")
	}
	return out, nil
}

// ChatStream is Chat, but calls onText with the reply so far each time the
// model writes more, so a caller can show it growing instead of waiting for
// the end. onText may be nil. The returned string is the complete reply.
func (c *Client) ChatStream(ctx context.Context, system, user string, onText func(string)) (string, error) {
	resp, err := c.open(ctx, system, user, true)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	// Ollama streams one JSON object per line.
	r := bufio.NewReader(resp.Body)
	var acc strings.Builder
	lastShown := ""

	for {
		line, readErr := r.ReadBytes('\n')
		if len(bytes.TrimSpace(line)) > 0 {
			var parsed chatResponse
			if err := json.Unmarshal(line, &parsed); err != nil {
				return "", fmt.Errorf("decoding ollama stream: %w", err)
			}
			if parsed.Error != "" {
				return "", fmt.Errorf("ollama error: %s", parsed.Error)
			}
			if parsed.Message.Content != "" {
				acc.WriteString(parsed.Message.Content)
				if shown := visible(acc.String()); onText != nil && shown != "" && shown != lastShown {
					lastShown = shown
					onText(shown)
				}
			}
			if parsed.Done {
				break
			}
		}
		if readErr != nil {
			if readErr == io.EOF {
				break
			}
			return "", fmt.Errorf("ollama stream interrupted: %w", readErr)
		}
	}

	out := answer(acc.String())
	if out == "" {
		return "", fmt.Errorf("ollama returned an empty reply")
	}
	return out, nil
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "…"
}