addrkit/internal/adapters/llm/rewriter.go

// Package llm implements the optional, disabled-by-default query rewriter.
//
// It talks to any OpenAI-compatible chat completions endpoint (for example the
// on-prem vLLM already used by AMS). Its contract is deliberately narrow: it
// may only propose alternative spellings of the user's address. It never
// returns the answer, and its output is validated by the deterministic
// resolver, so a misbehaving model can at worst waste a request — it cannot
// make the service answer with a wrong address.
package llm

import (
	"bytes"
	"context"
	_ "embed"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"strings"
	"time"
	"unicode/utf8"

	"git.totmin.ru/en2zmax/addrkit/internal/config"
)

// defaultSystemPrompt is embedded from a readable text file so the prompt can
// be reviewed and versioned like documentation. It can be overridden per
// deployment via llm.system_prompt or llm.system_prompt_file in the config.
//
//go:embed default_prompt.txt
var defaultSystemPrompt string

// Rewriter asks a model for normalised variants of an address.
type Rewriter struct {
	endpoint        string
	token           string
	model           string
	systemPrompt    string
	maxTokens       int
	disableThinking bool
	http            *http.Client
}

// New builds a rewriter from configuration. The URL is treated as a base; the
// OpenAI path is appended if it is not already present.
func New(cfg config.LLMConfig) *Rewriter {
	base := strings.TrimRight(cfg.URL, "/")
	if !strings.HasSuffix(base, "/chat/completions") {
		base += "/v1/chat/completions"
	}
	timeout := time.Duration(cfg.TimeoutSeconds) * time.Second
	if timeout <= 0 {
		timeout = 3 * time.Second
	}
	maxTokens := cfg.MaxTokens
	if maxTokens <= 0 {
		maxTokens = 512
	}
	prompt := strings.TrimSpace(cfg.SystemPrompt)
	if prompt == "" {
		prompt = strings.TrimSpace(defaultSystemPrompt)
	}
	return &Rewriter{
		endpoint:        base,
		token:           cfg.Token,
		model:           cfg.Model,
		systemPrompt:    prompt,
		maxTokens:       maxTokens,
		disableThinking: cfg.DisableThinking,
		http:            &http.Client{Timeout: timeout},
	}
}

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

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

type chatResponse struct {
	Choices []struct {
		Message chatMessage `json:"message"`
	} `json:"choices"`
}

// Rewrite returns up to three normalised variants of raw.
func (r *Rewriter) Rewrite(ctx context.Context, raw string) ([]string, error) {
	reqBody := chatRequest{
		Model: r.model,
		Messages: []chatMessage{
			{Role: "system", Content: r.systemPrompt},
			{Role: "user", Content: raw},
		},
		Temperature: 0,
		MaxTokens:   r.maxTokens,
	}
	if r.disableThinking {
		// Qwen3 and friends: skip chain-of-thought, we only need the JSON.
		reqBody.ChatTemplateKwargs = map[string]any{"enable_thinking": false}
	}
	payload, err := json.Marshal(reqBody)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.endpoint, bytes.NewReader(payload))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	if r.token != "" && r.token != "EMPTY" {
		req.Header.Set("Authorization", "Bearer "+r.token)
	}

	resp, err := r.http.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("llm status %d", resp.StatusCode)
	}

	var cr chatResponse
	// Bound the reply: a broken or hostile model must not be able to stream an
	// unbounded JSON body into the service.
	if err := json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&cr); err != nil {
		return nil, fmt.Errorf("decode llm response: %w", err)
	}
	if len(cr.Choices) == 0 {
		return nil, nil
	}
	return parseVariants(cr.Choices[0].Message.Content), nil
}

// parseVariants extracts the {"variants":[...]} object from a model reply that
// may be wrapped in reasoning or code fences. Thinking models are asked to
// skip reasoning, but we tolerate it anyway.
func parseVariants(content string) []string {
	content = stripThinking(content)
	raw := extractVariantsObject(content)
	if raw == "" {
		return nil
	}
	var parsed struct {
		Variants []string `json:"variants"`
	}
	if err := json.Unmarshal([]byte(raw), &parsed); err != nil {
		return nil
	}
	out := make([]string, 0, len(parsed.Variants))
	for _, v := range parsed.Variants {
		s := strings.TrimSpace(v)
		// A variant is a short address; the cap also keeps a hostile model
		// from making the resolver scan a megabyte-long token.
		if s == "" || utf8.RuneCountInString(s) > 512 {
			continue
		}
		out = append(out, s)
		if len(out) == 3 {
			break
		}
	}
	return out
}

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

// stripThinking removes  blocks, including an unterminated one left
// behind when the token budget is exhausted mid-reasoning.
func stripThinking(content string) string {
	content = reThinkBlock.ReplaceAllString(content, "")
	if i := strings.Index(content, "<think>"); i >= 0 {
		if j := strings.Index(content[i:], "</think>"); j < 0 {
			content = content[:i]
		}
	}
	return content
}

// extractVariantsObject returns the smallest {...} that contains "variants".
func extractVariantsObject(content string) string {
	idx := strings.Index(content, `"variants"`)
	if idx < 0 {
		return ""
	}
	start := strings.LastIndex(content[:idx], "{")
	if start < 0 {
		return ""
	}
	end := strings.Index(content[idx:], "}")
	if end < 0 {
		return ""
	}
	return content[start : idx+end+1]
}