compass_teams_nhs/internal/llm/omniroute.go

// Package llm is a small OpenAI-compatible client for OmniRoute. All model
// traffic goes through OmniRoute (AGENTS.md §10 rule 4); this package holds no
// model list and no credentials of its own.
package llm

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

// DefaultBaseURL is OmniRoute on the host that runs this service (.13).
const DefaultBaseURL = "http://127.0.0.1:20129/v1"

// DefaultModel is a cheap, reliable text model used for importance labels. It is
// overridable with SCHOOL_LLM_MODEL.
const DefaultModel = "openrouter/google/gemini-3.5-flash-lite"

// Client calls one OmniRoute chat-completions endpoint.
type Client struct {
	BaseURL string
	APIKey  string
	Model   string
	HTTP    *http.Client
}

// New builds a client. Empty baseURL and model fall back to the defaults.
func New(baseURL, apiKey, model string) *Client {
	if strings.TrimSpace(baseURL) == "" {
		baseURL = DefaultBaseURL
	}
	if strings.TrimSpace(model) == "" {
		model = DefaultModel
	}
	return &Client{
		BaseURL: strings.TrimRight(baseURL, "/"),
		APIKey:  apiKey,
		Model:   model,
		HTTP:    &http.Client{Timeout: 90 * time.Second},
	}
}

// Complete sends one system+user prompt and returns the assistant text. It
// accepts both a normal JSON reply and a streaming SSE reply, because some
// OmniRoute providers stream regardless of the stream flag.
func (c *Client) Complete(ctx context.Context, system, prompt string) (string, error) {
	if c.APIKey == "" {
		return "", fmt.Errorf("omniroute: no API key")
	}
	body, err := json.Marshal(map[string]any{
		"model": c.Model,
		"messages": []map[string]string{
			{"role": "system", "content": system},
			{"role": "user", "content": prompt},
		},
		"temperature": 0,
		"max_tokens":  4096,
		"stream":      false,
	})
	if err != nil {
		return "", fmt.Errorf("omniroute: encode request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return "", fmt.Errorf("omniroute: request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.APIKey)

	resp, err := c.HTTP.Do(req)
	if err != nil {
		return "", fmt.Errorf("omniroute: %w", err)
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
	if err != nil {
		return "", fmt.Errorf("omniroute: read response: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("omniroute: HTTP %d: %s", resp.StatusCode, truncate(string(raw), 400))
	}

	if strings.Contains(resp.Header.Get("Content-Type"), "text/event-stream") || bytes.HasPrefix(bytes.TrimSpace(raw), []byte("data:")) {
		return parseSSE(raw)
	}

	var parsed struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return "", fmt.Errorf("omniroute: decode response: %w", err)
	}
	if len(parsed.Choices) == 0 {
		return "", fmt.Errorf("omniroute: no choices in response: %s", truncate(string(raw), 400))
	}
	return parsed.Choices[0].Message.Content, nil
}

// parseSSE concatenates the delta content of an SSE chat-completions stream.
func parseSSE(raw []byte) (string, error) {
	var b strings.Builder
	scanner := bufio.NewScanner(bytes.NewReader(raw))
	scanner.Buffer(make([]byte, 0, 64*1024), 1<<20)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if payload == "" || payload == "[DONE]" {
			continue
		}
		var chunk struct {
			Choices []struct {
				Delta struct {
					Content string `json:"content"`
				} `json:"delta"`
			} `json:"choices"`
		}
		if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
			continue
		}
		for _, ch := range chunk.Choices {
			b.WriteString(ch.Delta.Content)
		}
	}
	if err := scanner.Err(); err != nil {
		return "", fmt.Errorf("omniroute: read stream: %w", err)
	}
	if b.Len() == 0 {
		return "", fmt.Errorf("omniroute: empty streamed response")
	}
	return b.String(), nil
}

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