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