tapir/internal/adapters/llm/client.go

// Package llm is a COPY of hyperguild/ingestion/internal/llm (ADR-004). Tapir
// owns this copy; it is not a module dependency on that repo. The package is
// stdlib-only. Client calls an OpenAI-compatible chat completions endpoint;
// Router (router.go) implements the local-Primary -> BYO-Fallback pattern Tapir
// needs for docs/use-cases/ai_routing.feature.
package llm

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

// defaultMaxTokens is sent on every request. Tapir CHANGES this from the
// hyperguild copy (ADR-004 says change the copy, not the upstream): thinking
// models (qwen3, deepseek-r1) spend their budget on reasoning and return EMPTY
// content when max_tokens is unset or too low. A generous ceiling leaves room
// for both the reasoning trace and the actual summary. See
// docs/homelab-integration.md.
const defaultMaxTokens = 8192

// Client calls an OpenAI-compatible chat completions endpoint.
type Client struct {
	baseURL    string
	apiKey     string
	model      string
	maxTokens  int
	httpClient *http.Client
	usageHook  func(model string, prompt, completion int)
}

// Option configures a Client at construction. Variadic so the existing 4-arg
// call sites stay valid as new knobs are added.
type Option func(*Client)

// WithMaxTokens overrides the per-request completion budget. The summarizer uses
// this to cap completion for small-context models (e.g. koala/phi4-mini, 8k):
// with the default 8192 budget, prompt + max_tokens overflows an 8k context and
// the gateway returns HTTP 400. A non-positive n is ignored (keeps the default).
func WithMaxTokens(n int) Option {
	return func(c *Client) {
		if n > 0 {
			c.maxTokens = n
		}
	}
}

// WithUsageHook registers a callback fired after a successful completion with the
// model and the prompt/completion token counts from the response usage block. It
// keeps this copied, stdlib-only package (ADR-004) decoupled from metrics: the
// caller wires it to internal/metrics, the client imports nothing. nil is ignored.
func WithUsageHook(fn func(model string, prompt, completion int)) Option {
	return func(c *Client) { c.usageHook = fn }
}

// New constructs a Client.
func New(baseURL, apiKey, model string, timeout time.Duration, opts ...Option) *Client {
	c := &Client{
		baseURL:    strings.TrimRight(baseURL, "/"),
		apiKey:     apiKey,
		model:      model,
		maxTokens:  defaultMaxTokens,
		httpClient: &http.Client{Timeout: timeout},
	}
	for _, opt := range opts {
		opt(c)
	}
	return c
}

type chatRequest struct {
	Model       string    `json:"model"`
	Messages    []message `json:"messages"`
	Temperature float64   `json:"temperature"`
	MaxTokens   int       `json:"max_tokens,omitempty"`
}

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

type chatResponse struct {
	Choices []struct {
		Message message `json:"message"`
	} `json:"choices"`
	Usage struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
	} `json:"usage"`
}

// Complete sends a system + user message and returns the assistant's reply.
// Retries once on HTTP 429 using Retry-After header or 5s backoff.
func (c *Client) Complete(ctx context.Context, system, user string) (string, error) {
	body := chatRequest{
		Model: c.model,
		Messages: []message{
			{Role: "system", Content: system},
			{Role: "user", Content: user},
		},
		Temperature: 0.2,
		MaxTokens:   c.maxTokens,
	}
	b, err := json.Marshal(body)
	if err != nil {
		return "", fmt.Errorf("marshal request: %w", err)
	}

	do := func() (*http.Response, error) {
		req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/chat/completions", bytes.NewReader(b))
		if err != nil {
			return nil, fmt.Errorf("build request: %w", err)
		}
		req.Header.Set("Content-Type", "application/json")
		if c.apiKey != "" {
			req.Header.Set("Authorization", "Bearer "+c.apiKey)
		}
		return c.httpClient.Do(req)
	}

	resp, err := do()
	if err != nil {
		return "", fmt.Errorf("call LLM: %w", err)
	}

	if resp.StatusCode == http.StatusTooManyRequests {
		_ = resp.Body.Close()
		wait := 5 * time.Second
		if ra := resp.Header.Get("Retry-After"); ra != "" {
			if secs, err := strconv.Atoi(ra); err == nil {
				wait = time.Duration(secs) * time.Second
			}
		}
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case <-time.After(wait):
		}
		resp, err = do()
		if err != nil {
			return "", fmt.Errorf("retry LLM call: %w", err)
		}
	}
	defer resp.Body.Close() //nolint:errcheck

	out, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("read response: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("LLM returned %d: %s", resp.StatusCode, out)
	}

	var cr chatResponse
	if err := json.Unmarshal(out, &cr); err != nil {
		return "", fmt.Errorf("parse response: %w", err)
	}
	if len(cr.Choices) == 0 {
		return "", fmt.Errorf("LLM returned no choices")
	}
	if c.usageHook != nil {
		c.usageHook(c.model, cr.Usage.PromptTokens, cr.Usage.CompletionTokens)
	}
	return cr.Choices[0].Message.Content, nil
}