| 1 | package retrieval |
| 2 | |
| 3 | import ( |
| 4 | "fmt" |
| 5 | "math" |
| 6 | "strings" |
| 7 | "unicode" |
| 8 | "unicode/utf8" |
| 9 | ) |
| 10 | |
| 11 | // Tokens lowercases Latin words and splits CJK runs into overlapping bigrams |
| 12 | // (a lone CJK rune stays a single term) — the standard CJK indexing unit |
| 13 | // (Lucene's CJKAnalyzer, SQLite FTS): a bigram only matches a real two-rune |
| 14 | // subsequence where per-rune unigrams matched scattered common characters. |
| 15 | // Intentionally a local, dependency-free approximation of FTS matching. |
| 16 | func Tokens(s string) []string { |
| 17 | var out []string |
| 18 | var b strings.Builder |
| 19 | var prev rune |
| 20 | cjkRun := 0 |
| 21 | flush := func() { |
| 22 | if b.Len() == 0 { |
| 23 | return |
| 24 | } |
| 25 | out = append(out, b.String()) |
| 26 | b.Reset() |
| 27 | } |
| 28 | endCJK := func() { |
| 29 | if cjkRun == 1 { |
| 30 | out = append(out, string(prev)) |
| 31 | } |
| 32 | cjkRun = 0 |
| 33 | } |
| 34 | for _, r := range s { |
| 35 | switch { |
| 36 | case isCJK(r): |
| 37 | flush() |
| 38 | if cjkRun > 0 { |
| 39 | out = append(out, string([]rune{prev, r})) |
| 40 | } |
| 41 | prev = r |
| 42 | cjkRun++ |
| 43 | case unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_': |
| 44 | endCJK() |
| 45 | b.WriteRune(unicode.ToLower(r)) |
| 46 | default: |
| 47 | flush() |
| 48 | endCJK() |
| 49 | } |
| 50 | } |
| 51 | flush() |
| 52 | endCJK() |
| 53 | return out |
| 54 | } |
| 55 | |
| 56 | func isCJK(r rune) bool { |
| 57 | return unicode.In(r, unicode.Han, unicode.Hiragana, unicode.Katakana, unicode.Hangul) |
| 58 | } |
| 59 | |
| 60 | // Unique returns terms in first-seen order. |
| 61 | func Unique(in []string) []string { |
| 62 | seen := map[string]bool{} |
| 63 | out := make([]string, 0, len(in)) |
| 64 | for _, s := range in { |
| 65 | if s == "" || seen[s] { |
| 66 | continue |
| 67 | } |
| 68 | seen[s] = true |
| 69 | out = append(out, s) |
| 70 | } |
| 71 | return out |
| 72 | } |
| 73 | |
| 74 | // Counts returns a term-frequency map. |
| 75 | func Counts(terms []string) map[string]int { |
| 76 | counts := map[string]int{} |
| 77 | for _, term := range terms { |
| 78 | counts[term]++ |
| 79 | } |
| 80 | return counts |
| 81 | } |
| 82 | |
| 83 | // BM25Score scores a document against query terms. |
| 84 | func BM25Score(counts map[string]int, length int, queryTerms []string, df map[string]int, totalDocs int, avgLen float64) float64 { |
| 85 | const ( |
| 86 | k1 = 1.2 |
| 87 | b = 0.75 |
| 88 | ) |
| 89 | if length <= 0 || totalDocs <= 0 { |
| 90 | return 0 |
| 91 | } |
| 92 | if avgLen <= 0 { |
| 93 | avgLen = 1 |
| 94 | } |
| 95 | var score float64 |
| 96 | docLen := float64(length) |
| 97 | for _, term := range queryTerms { |
| 98 | tf := counts[term] |
| 99 | if tf == 0 { |
| 100 | continue |
| 101 | } |
| 102 | termDF := df[term] |
| 103 | if termDF == 0 { |
| 104 | continue |
| 105 | } |
| 106 | idf := math.Log(1 + (float64(totalDocs)-float64(termDF)+0.5)/(float64(termDF)+0.5)) |
| 107 | freq := float64(tf) |
| 108 | score += idf * (freq * (k1 + 1)) / (freq + k1*(1-b+b*docLen/avgLen)) |
| 109 | } |
| 110 | return score |
| 111 | } |
| 112 | |
| 113 | // DocumentFrequency counts how many documents contain each term. |
| 114 | func DocumentFrequency(docs []map[string]int) map[string]int { |
| 115 | df := map[string]int{} |
| 116 | for _, counts := range docs { |
| 117 | for term := range counts { |
| 118 | df[term]++ |
| 119 | } |
| 120 | } |
| 121 | return df |
| 122 | } |
| 123 | |
| 124 | // KeepTopRelativeScore keeps the best item and drops trailing items whose score |
| 125 | // falls below ratio * topScore. Callers must pass items already sorted best |
| 126 | // first. This mirrors SQLite FTS/BM25 search UIs that over-fetch, then trim |
| 127 | // common-word-only noise without imposing an absolute score threshold. |
| 128 | func KeepTopRelativeScore[T any](items []T, ratio float64, score func(T) float64) []T { |
| 129 | if len(items) == 0 || ratio <= 0 { |
| 130 | return items |
| 131 | } |
| 132 | top := score(items[0]) |
| 133 | if top <= 0 { |
| 134 | return items |
| 135 | } |
| 136 | cutoff := top * ratio |
| 137 | out := items[:0] |
| 138 | for i, item := range items { |
| 139 | if i == 0 || score(item) >= cutoff { |
| 140 | out = append(out, item) |
| 141 | } |
| 142 | } |
| 143 | return out |
| 144 | } |
| 145 | |
| 146 | // QueryTerms normalizes a search string and reports an error when nothing |
| 147 | // searchable remains. |
| 148 | func QueryTerms(query string) ([]string, error) { |
| 149 | terms := Unique(Tokens(strings.TrimSpace(query))) |
| 150 | if len(terms) == 0 { |
| 151 | return nil, fmt.Errorf("query must contain at least one letter or number") |
| 152 | } |
| 153 | return terms, nil |
| 154 | } |
| 155 | |
| 156 | // MakeSnippet returns a whitespace-compacted excerpt centered near the query. |
| 157 | func MakeSnippet(text, query string, terms []string, maxRunes int) string { |
| 158 | text = CompactWhitespace(text) |
| 159 | if maxRunes <= 0 || utf8.RuneCountInString(text) <= maxRunes { |
| 160 | return text |
| 161 | } |
| 162 | lower := strings.ToLower(text) |
| 163 | query = strings.ToLower(strings.TrimSpace(query)) |
| 164 | idx := -1 |
| 165 | if query != "" { |
| 166 | idx = strings.Index(lower, query) |
| 167 | } |
| 168 | if idx < 0 { |
| 169 | for _, term := range terms { |
| 170 | runes := []rune(term) |
| 171 | if len(runes) == 1 && !isCJK(runes[0]) { |
| 172 | continue |
| 173 | } |
| 174 | if i := strings.Index(lower, term); i >= 0 { |
| 175 | idx = i |
| 176 | break |
| 177 | } |
| 178 | } |
| 179 | } |
| 180 | if idx < 0 { |
| 181 | idx = 0 |
| 182 | } |
| 183 | return snippetAround(text, idx, maxRunes) |
| 184 | } |
| 185 | |
| 186 | func snippetAround(text string, byteIdx, maxRunes int) string { |
| 187 | if byteIdx < 0 { |
| 188 | byteIdx = 0 |
| 189 | } |
| 190 | if byteIdx > len(text) { |
| 191 | byteIdx = len(text) |
| 192 | } |
| 193 | for byteIdx > 0 && byteIdx < len(text) && !utf8.RuneStart(text[byteIdx]) { |
| 194 | byteIdx-- |
| 195 | } |
| 196 | runes := []rune(text) |
| 197 | pos := utf8.RuneCountInString(text[:byteIdx]) |
| 198 | start := max(pos-maxRunes/2, 0) |
| 199 | end := start + maxRunes |
| 200 | if end > len(runes) { |
| 201 | end = len(runes) |
| 202 | start = max(end-maxRunes, 0) |
| 203 | } |
| 204 | prefix := "" |
| 205 | suffix := "" |
| 206 | if start > 0 { |
| 207 | prefix = "..." |
| 208 | } |
| 209 | if end < len(runes) { |
| 210 | suffix = "..." |
| 211 | } |
| 212 | return prefix + string(runes[start:end]) + suffix |
| 213 | } |
| 214 | |
| 215 | // CompactWhitespace collapses runs of whitespace into one ASCII space. |
| 216 | func CompactWhitespace(s string) string { |
| 217 | return strings.Join(strings.Fields(s), " ") |
| 218 | } |
| 219 |