addrkit: strict ordinal matching, LLM-variant guard, externalized prompt

- match: require candidate ordinal to equal the query ordinal in fuzzy search (deterministic typo recall 73% -> 80%, no wrong houses) - resolver: compatibleVariant guard — LLM may only fix street spelling, keeping the same words, ordinal and house unit; damerau distance for transpositions - llm: prompt moved out of code (go:embed default_prompt.txt) with llm.system_prompt / llm.system_prompt_file overrides; guided JSON parsing tolerant of thinking blocks - address: recognize ordinals appearing after the street name - tests: first coverage for the LLM adapter and config prompt loading; README Measured (sample 200): deterministic 94.7-95.8% exact / 0 wrong across 4 seeds; LLM+guard typo 90.7-89.8% / 0 wrong. LLM fallback latency ~1.2-1.9s.
2026-10-09 21:40:46 +07:00
parent ba4705fc4f
commit f1537f1d6b
17 changed files with 590 additions and 48 deletions
@@ -19,7 +19,7 @@ import (
 func testResolver(t *testing.T) *match.Resolver {
 	t.Helper()
 	var entries []address.Entry
-	for _, a := range []string{"Станционная улица 1/1", "Искитимская улица 5"} {
+	for _, a := range []string{"Станционная улица 1/1", "2-я Станционная улица 1/1", "Искитимская улица 5"} {
 		p := address.Parse(a)
 		entries = append(entries, address.Entry{
 			ID: a, Address: a, District: "Тестовый",
@@ -68,7 +68,7 @@ func TestResolveEndpoint(t *testing.T) {
 	if err := json.NewDecoder(resp.Body).Decode(&res); err != nil {
 		t.Fatalf("decode: %v", err)
 	}
-	if res.Status != address.StatusFound || res.Address != "Станционная улица 1/1" {
+	if res.Status != address.StatusFound || res.Address != "2-я Станционная улица 1/1" {
 		t.Fatalf("unexpected result: %+v", res)
 	}
 }
@@ -0,0 +1,11 @@
+Ты — нормализатор российских адресов. На входе адрес, который человек написал как угодно: с опечатками, перестановками букв, сокращениями и лишними словами.
+Верни СТРОГО JSON: {"variants":["...","..."]} — до 3 вариантов ИСПРАВЛЕННОЙ записи того же адреса, от самого вероятного к менее вероятному.
+Формат варианта: "<Название улицы> <дом>[/<дробь>][ <литера>][ к<корпус>][ стр.<строение>]".
+Правила:
+- исправляй опечатки и перестановки букв в названии улицы;
+- раскрывай сокращения: ул.→улица, пер.→переулок, пр-т→проспект, ш.→шоссе, д.→дом;
+- восстанавливай порядковые: "2 я", "вторая" → "2-я";
+- убирай "город", "район", "квартира" и прочий мусор;
+- НЕ заменяй улицу на другую и НЕ добавляй дом, которого нет во входе: меняй только написание;
+- если адрес уже верный — верни его без изменений.
+Отвечай только JSON, без пояснений.
@@ -11,21 +11,33 @@ package llm
 import (
 	"bytes"
 	"context"
+	_ "embed"
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"regexp"
 	"strings"
 	"time"

 	"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
-	http     *http.Client
+	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
@@ -39,24 +51,31 @@ func New(cfg config.LLMConfig) *Rewriter {
 	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,
-		http:     &http.Client{Timeout: timeout},
+		endpoint:        base,
+		token:           cfg.Token,
+		model:           cfg.Model,
+		systemPrompt:    prompt,
+		maxTokens:       maxTokens,
+		disableThinking: cfg.DisableThinking,
+		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"`
+	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 {
@@ -72,15 +91,20 @@ type chatResponse struct {

 // 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{
+	reqBody := chatRequest{
 		Model: r.model,
 		Messages: []chatMessage{
-			{Role: "system", Content: systemPrompt},
+			{Role: "system", Content: r.systemPrompt},
 			{Role: "user", Content: raw},
 		},
 		Temperature: 0,
-		MaxTokens:   256,
-	})
+		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
 	}
@@ -114,17 +138,18 @@ func (r *Rewriter) Rewrite(ctx context.Context, raw string) ([]string, error) {
 }

 // parseVariants extracts the {"variants":[...]} object from a model reply that
-// may be wrapped in prose or code fences.
+// 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 {
-	start := strings.Index(content, "{")
-	end := strings.LastIndex(content, "}")
-	if start < 0 || end <= start {
+	content = stripThinking(content)
+	raw := extractVariantsObject(content)
+	if raw == "" {
 		return nil
 	}
 	var parsed struct {
 		Variants []string `json:"variants"`
 	}
-	if err := json.Unmarshal([]byte(content[start:end+1]), &parsed); err != nil {
+	if err := json.Unmarshal([]byte(raw), &parsed); err != nil {
 		return nil
 	}
 	out := make([]string, 0, len(parsed.Variants))
@@ -135,3 +160,34 @@ func parseVariants(content string) []string {
 	}
 	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]
+}
@@ -0,0 +1,88 @@
+package llm
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"git.totmin.ru/en2zmax/addrkit/internal/config"
+)
+
+// chatServer returns a test server that echoes a fixed content and captures the
+// request body it received.
+func chatServer(t *testing.T, content string) (*httptest.Server, *chatRequest) {
+	t.Helper()
+	var got chatRequest
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_ = json.NewDecoder(r.Body).Decode(&got)
+		w.Header().Set("Content-Type", "application/json")
+		resp := map[string]any{"choices": []any{map[string]any{"message": map[string]any{"content": content}}}}
+		_ = json.NewEncoder(w).Encode(resp)
+	}))
+	t.Cleanup(ts.Close)
+	return ts, &got
+}
+
+func TestRewriteUsesConfiguredPrompt(t *testing.T) {
+	ts, got := chatServer(t, `{"variants":["А"]}`)
+	r := New(config.LLMConfig{URL: ts.URL, Model: "m", SystemPrompt: "CUSTOM PROMPT"})
+
+	out, err := r.Rewrite(context.Background(), "вход")
+	if err != nil {
+		t.Fatalf("Rewrite: %v", err)
+	}
+	if len(got.Messages) == 0 || got.Messages[0].Content != "CUSTOM PROMPT" {
+		t.Errorf("system prompt not applied: %+v", got.Messages)
+	}
+	if len(out) != 1 || out[0] != "А" {
+		t.Errorf("variants = %v", out)
+	}
+}
+
+func TestRewriteFallsBackToEmbeddedPrompt(t *testing.T) {
+	ts, got := chatServer(t, `{"variants":[]}`)
+	r := New(config.LLMConfig{URL: ts.URL, Model: "m"})
+
+	if _, err := r.Rewrite(context.Background(), "x"); err != nil {
+		t.Fatalf("Rewrite: %v", err)
+	}
+	if len(got.Messages) == 0 || !strings.Contains(got.Messages[0].Content, "нормализатор") {
+		t.Errorf("embedded default prompt not used: %q", got.Messages[0].Content)
+	}
+}
+
+func TestRewriteDisablesThinking(t *testing.T) {
+	ts, got := chatServer(t, `{"variants":[]}`)
+	r := New(config.LLMConfig{URL: ts.URL, Model: "m", DisableThinking: true})
+
+	if _, err := r.Rewrite(context.Background(), "x"); err != nil {
+		t.Fatalf("Rewrite: %v", err)
+	}
+	if v, ok := got.ChatTemplateKwargs["enable_thinking"]; !ok || v != false {
+		t.Errorf("enable_thinking not disabled: %+v", got.ChatTemplateKwargs)
+	}
+}
+
+func TestParseVariants(t *testing.T) {
+	cases := []struct {
+		name    string
+		content string
+		want    int
+	}{
+		{"clean", `{"variants":["a","b"]}`, 2},
+		{"with thinking", "<think>hmm</think>\n{\"variants\":[\"a\"]}", 1},
+		{"unterminated thinking", "<think>hmm {\"variants\":[\"a\"]}", 0},
+		{"prose around", "Here you go: {\"variants\":[\"a\",\"b\",\"c\"]} done", 3},
+		{"garbage", "no json here", 0},
+	}
+	for _, tt := range cases {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := parseVariants(tt.content); len(got) != tt.want {
+				t.Errorf("parseVariants(%q) = %v, want %d items", tt.content, got, tt.want)
+			}
+		})
+	}
+}