| 1 | package retrieval |
| 2 | |
| 3 | import ( |
| 4 | "strings" |
| 5 | "testing" |
| 6 | "unicode/utf8" |
| 7 | ) |
| 8 | |
| 9 | func TestTokensHandlesLatinAndCJK(t *testing.T) { |
| 10 | for _, tc := range []struct { |
| 11 | in string |
| 12 | want []string |
| 13 | }{ |
| 14 | {"BM25 检索 cache-first", []string{"bm25", "检索", "cache", "first"}}, |
| 15 | {"数据库迁移", []string{"数据", "据库", "库迁", "迁移"}}, |
| 16 | {"库", []string{"库"}}, |
| 17 | {"用pnpm装依赖", []string{"用", "pnpm", "装依", "依赖"}}, |
| 18 | } { |
| 19 | got := Tokens(tc.in) |
| 20 | if strings.Join(got, ",") != strings.Join(tc.want, ",") { |
| 21 | t.Fatalf("Tokens(%q) = %#v, want %#v", tc.in, got, tc.want) |
| 22 | } |
| 23 | } |
| 24 | } |
| 25 | |
| 26 | func TestBM25ScoreRanksMatchingDocument(t *testing.T) { |
| 27 | query := Unique(Tokens("prompt cache")) |
| 28 | doc1 := Counts(Tokens("prompt cache cache stability")) |
| 29 | doc2 := Counts(Tokens("dashboard colors")) |
| 30 | df := DocumentFrequency([]map[string]int{doc1, doc2}) |
| 31 | score1 := BM25Score(doc1, 4, query, df, 2, 3) |
| 32 | score2 := BM25Score(doc2, 2, query, df, 2, 3) |
| 33 | if score1 <= score2 { |
| 34 | t.Fatalf("matching score %.3f should exceed unrelated score %.3f", score1, score2) |
| 35 | } |
| 36 | } |
| 37 | |
| 38 | func TestKeepTopRelativeScoreKeepsTopAndDropsWeakTail(t *testing.T) { |
| 39 | items := []struct { |
| 40 | name string |
| 41 | score float64 |
| 42 | }{ |
| 43 | {name: "top", score: 10}, |
| 44 | {name: "near", score: 2}, |
| 45 | {name: "noise", score: 1.4}, |
| 46 | {name: "zero", score: 0}, |
| 47 | } |
| 48 | got := KeepTopRelativeScore(items, 0.15, func(item struct { |
| 49 | name string |
| 50 | score float64 |
| 51 | }) float64 { |
| 52 | return item.score |
| 53 | }) |
| 54 | if len(got) != 2 || got[0].name != "top" || got[1].name != "near" { |
| 55 | t.Fatalf("KeepTopRelativeScore() = %#v, want top and near", got) |
| 56 | } |
| 57 | } |
| 58 | |
| 59 | func TestMakeSnippetHandlesMultibyteBoundary(t *testing.T) { |
| 60 | text := strings.Repeat("前缀", 80) + "稳定结论 synthesis cache " + strings.Repeat("后缀", 80) |
| 61 | out := MakeSnippet(text, "synthesis cache", QueryTermsForTest(t, "synthesis cache"), 60) |
| 62 | if !strings.Contains(out, "synthesis cache") { |
| 63 | t.Fatalf("snippet missing query: %q", out) |
| 64 | } |
| 65 | if strings.ContainsRune(out, utf8.RuneError) { |
| 66 | t.Fatalf("snippet contains replacement rune: %q", out) |
| 67 | } |
| 68 | } |
| 69 | |
| 70 | func QueryTermsForTest(t *testing.T, query string) []string { |
| 71 | t.Helper() |
| 72 | terms, err := QueryTerms(query) |
| 73 | if err != nil { |
| 74 | t.Fatal(err) |
| 75 | } |
| 76 | return terms |
| 77 | } |
| 78 |