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