Maksim Totmin
17ff5c591f
- remove the events/outages module, ArcGIS adapter, portal config and event statuses - add a configurable city dictionary/parser and make City part of StreetKey - seed: read city/city_name, namespace IDs by city, keep structures distinct - fix reliability: schema v2 with fail-closed load, deterministic ordering, input limits, config validation, atomic DB publish, LLM response caps - add read-only shadow mode comparing the old ES path with addrkit - breaking: registry.db must be re-seeded; config has no portal section
103 lines
3.1 KiB
Go
103 lines
3.1 KiB
Go
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)
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestParseVariantsCapsCountAndLength(t *testing.T) {
|
||
long := strings.Repeat("x", 600)
|
||
content := `{"variants":["` + long + `","a","b","c","d"]}`
|
||
out := parseVariants(content)
|
||
if len(out) != 3 {
|
||
t.Fatalf("variants = %v, want 3", out)
|
||
}
|
||
for _, v := range out {
|
||
if len([]rune(v)) > 512 {
|
||
t.Fatalf("oversized variant accepted: %d runes", len([]rune(v)))
|
||
}
|
||
}
|
||
}
|