petal/internal/llm/client.go

// Package llm wraps the local inference endpoint behind a backend-agnostic
// interface. Two concrete backends — vLLM (OpenAI-compatible) and Ollama
// (native) — implement LLMClient; the rest of the app (the grammar checkpoint,
// Ask Petal) calls the interface only and never knows which one is active.
package llm

import (
	"context"
	"time"

	"gitea.parodia.dev/drwily/petal/internal/config"
)

// Message is one chat turn sent to the model.
type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// CompletionRequest is the backend-neutral request shape. Each backend maps
// these fields onto its own wire format (see the parameter cheatsheet in the
// spec). Zero-valued sampling fields fall back to backend defaults.
type CompletionRequest struct {
	Messages          []Message
	MaxTokens         int
	Temperature       float64
	RepetitionPenalty float64
	TopP              float64
	Stop              []string
	Stream            bool
}

// LLMClient is the single surface handler code depends on.
type LLMClient interface {
	// Complete returns the full response in one shot (used for the checkpoint
	// JSON, which we want whole before parsing).
	Complete(ctx context.Context, req CompletionRequest) (string, error)
	// Stream returns a channel of text chunks (used for Ask Petal SSE). The
	// channel closes when generation ends; on a mid-stream error it closes
	// early and the error is reported via the returned error of a future call.
	Stream(ctx context.Context, req CompletionRequest) (<-chan string, error)
}

// Truncation caps. These guard checkpoint latency (prefill time scales with
// input length), not the model window — Qwen 3.5 ships 256K, far larger.
const (
	// maxDocChars ~ 10,000 tokens at ~4 chars/token. Grammar checkpoint only;
	// the voice pass sends the whole document uncut.
	maxDocChars = 40000
	// maxHistoryMsgs caps the rolling Ask Petal history (5 turns); oldest pairs
	// drop first.
	maxHistoryMsgs = 10
)

// TruncateDoc keeps the recent end of a document — the user is actively writing
// there — when it exceeds the grammar-checkpoint latency cap. Most documents
// fit uncut. The voice pass deliberately does NOT call this.
func TruncateDoc(contentText string) string {
	if len(contentText) > maxDocChars {
		return contentText[len(contentText)-maxDocChars:]
	}
	return contentText
}

// TrimHistory keeps the most recent maxHistoryMsgs messages, dropping the
// oldest first so a long Ask Petal chat stays within budget.
func TrimHistory(msgs []Message) []Message {
	if len(msgs) > maxHistoryMsgs {
		return msgs[len(msgs)-maxHistoryMsgs:]
	}
	return msgs
}

// NewLLMClient selects a backend from config. LLM_CHAT_MODEL falls back to
// LLM_MODEL here, in the factory, so handlers never deal with the fallback.
func NewLLMClient(cfg *config.Config) LLMClient {
	chatModel := cfg.LLMChatModel
	if chatModel == "" {
		chatModel = cfg.LLMModel
	}
	base := backend{
		endpoint:        cfg.LLMEndpoint,
		checkpointModel: cfg.LLMModel,
		chatModel:       chatModel,
		timeout:         cfg.LLMTimeout,
	}
	switch cfg.LLMBackend {
	case "ollama":
		return &OllamaClient{base}
	default: // "vllm"
		return &VLLMClient{base}
	}
}

// backend holds the fields shared by both concrete clients.
type backend struct {
	endpoint        string
	checkpointModel string // LLM_MODEL — small/fast checkpoint model
	chatModel       string // LLM_CHAT_MODEL (or LLM_MODEL) — Ask Petal model
	timeout         time.Duration
}

// model picks the right model id for a request. Streaming requests are Ask
// Petal (chat); one-shot Complete calls are the grammar checkpoint.
func (b backend) model(req CompletionRequest) string {
	if req.Stream {
		return b.chatModel
	}
	return b.checkpointModel
}