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"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

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

// Rewriter asks a model for normalised variants of an address.
type Rewriter struct {
	endpoint string
	token    string
	model    string
	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
	}
	return &Rewriter{
		endpoint: base,
		token:    cfg.Token,
		model:    cfg.Model,
		http:     &http.Client{Timeout: timeout},
	}
}

const systemPrompt = `Ты — нормализатор адресов. Тебе дают адрес, написанный как угодно.
Верни ТОЛЬКО JSON вида {"variants":["...","..."]} — до 3 вариантов записи того же адреса
в формате "<название улицы> <дом>[/<дробь>][ <литера>][ к<корпус>][ стр.<строение>]".
Не выдумывай улицы и дома, которых нет во входе. Если не уверен — верни пустой список.`

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

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) {
	payload, err := json.Marshal(chatRequest{
		Model: r.model,
		Messages: []chatMessage{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: raw},
		},
		Temperature: 0,
		MaxTokens:   256,
	})
	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
	if err := json.NewDecoder(resp.Body).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 prose or code fences.
func parseVariants(content string) []string {
	start := strings.Index(content, "{")
	end := strings.LastIndex(content, "}")
	if start < 0 || end <= start {
		return nil
	}
	var parsed struct {
		Variants []string `json:"variants"`
	}
	if err := json.Unmarshal([]byte(content[start:end+1]), &parsed); err != nil {
		return nil
	}
	out := make([]string, 0, len(parsed.Variants))
	for _, v := range parsed.Variants {
		if s := strings.TrimSpace(v); s != "" {
			out = append(out, s)
		}
	}
	return out
}