| 1 | // Package index is memoria's search index. |
| 2 | // |
| 3 | // The index lives in memory and nowhere else. That is not a performance |
| 4 | // choice: FTS5 stores its term postings in plaintext, so an index on disk |
| 5 | // would put readable passages in the replicated database file and break |
| 6 | // the one rule. The sealed chunk rows are the durable copy; this is |
| 7 | // rebuilt from them on wake. |
| 8 | // |
| 9 | // The spec's §7.3 measures a rebuild at 19s for 90k chunks and rules |
| 10 | // "snapshot, never rebuild" above that. Slice 0 rebuilds, because at a |
| 11 | // few hundred documents it is milliseconds and a snapshot format is a |
| 12 | // thing to get wrong. The trigger to add one is written in §12. |
| 13 | package index |
| 14 | |
| 15 | import ( |
| 16 | "crypto/rand" |
| 17 | "database/sql" |
| 18 | "encoding/hex" |
| 19 | "fmt" |
| 20 | "strings" |
| 21 | "sync" |
| 22 | |
| 23 | _ "modernc.org/sqlite" |
| 24 | ) |
| 25 | |
| 26 | // Hit is one scored passage. It carries ids and a snippet, never the |
| 27 | // whole chunk: the caller re-reads the sealed row if it needs the body. |
| 28 | type Hit struct { |
| 29 | ChunkID int64 |
| 30 | DocumentID int64 |
| 31 | Ordinal int |
| 32 | Heading string |
| 33 | Snippet string |
| 34 | // Score is BM25, negated so larger is better. FTS5 returns bm25() |
| 35 | // with lower meaning more relevant, which reads backwards everywhere |
| 36 | // else in the codebase. |
| 37 | Score float64 |
| 38 | } |
| 39 | |
| 40 | // Index is an in-memory FTS5 index. Safe for concurrent use. |
| 41 | type Index struct { |
| 42 | mu sync.RWMutex |
| 43 | db *sql.DB |
| 44 | } |
| 45 | |
| 46 | // Open builds an empty index. The DSN is in-memory with a shared cache so |
| 47 | // the single handle is stable; nothing is written to any file. |
| 48 | func Open() (*Index, error) { |
| 49 | // `cache=shared` keeps one database alive across pooled connections. |
| 50 | // Without it each connection gets its own empty database and searches |
| 51 | // return nothing — a plausible empty list, which is the failure shape |
| 52 | // this codebase treats as the worst available, so it is worth the |
| 53 | // explicit DSN. |
| 54 | // |
| 55 | // The name must be UNIQUE PER INDEX, and that is the load-bearing |
| 56 | // half. A shared cache is keyed by name, so a fixed name makes every |
| 57 | // Index in the process the same database. Two apps in one process |
| 58 | // then share one index: the second Open fails with "table chunks |
| 59 | // already exists" if you are lucky, and if the schema ever matched it |
| 60 | // would instead merge two teams' passages into one searchable |
| 61 | // corpus — the cross-team leak §6.3 exists to make impossible. Found |
| 62 | // by the test suite building two apps in one process. |
| 63 | name, err := randomName() |
| 64 | if err != nil { |
| 65 | return nil, err |
| 66 | } |
| 67 | db, err := sql.Open("sqlite", "file:"+name+"?mode=memory&cache=shared") |
| 68 | if err != nil { |
| 69 | return nil, err |
| 70 | } |
| 71 | // One connection only. The shared-cache in-memory database is |
| 72 | // process-local and cheap; a pool buys nothing and makes the |
| 73 | // lifetime harder to reason about. |
| 74 | db.SetMaxOpenConns(1) |
| 75 | if _, err := db.Exec(` |
| 76 | CREATE VIRTUAL TABLE chunks USING fts5( |
| 77 | heading, |
| 78 | body, |
| 79 | -- title is the document's own title, repeated into every one |
| 80 | -- of its chunks. Without it a document is unfindable by the |
| 81 | -- words in its name: a pasted "Passport numbers" whose body |
| 82 | -- is numbers matched nothing for "passport", and the page |
| 83 | -- reported that as a real absence. |
| 84 | -- |
| 85 | -- It goes AFTER body deliberately. snippet() and bm25() |
| 86 | -- address columns by declared position, so putting it first |
| 87 | -- would move body from 1 to 2 and silently change which |
| 88 | -- column gets highlighted. path goes after title for the |
| 89 | -- same reason. |
| 90 | title, |
| 91 | -- path is the document's ancestry in its SOURCE — "Handbook |
| 92 | -- › Onboarding" for a Notion page nested two deep. Distinct |
| 93 | -- from title (this page's own name) and from heading (where |
| 94 | -- in this page), and weighted below both: see PathWeight. |
| 95 | path, |
| 96 | chunk_id UNINDEXED, |
| 97 | document_id UNINDEXED, |
| 98 | ordinal UNINDEXED, |
| 99 | tokenize = 'porter unicode61' |
| 100 | );`); err != nil { |
| 101 | db.Close() |
| 102 | return nil, fmt.Errorf("index: create fts5 table: %w", err) |
| 103 | } |
| 104 | return &Index{db: db}, nil |
| 105 | } |
| 106 | |
| 107 | // Close releases the index. |
| 108 | func (i *Index) Close() error { |
| 109 | i.mu.Lock() |
| 110 | defer i.mu.Unlock() |
| 111 | return i.db.Close() |
| 112 | } |
| 113 | |
| 114 | // Entry is one chunk to index. Text is plaintext: it has been opened from |
| 115 | // its sealed row and must not outlive this process. |
| 116 | type Entry struct { |
| 117 | ChunkID int64 |
| 118 | DocumentID int64 |
| 119 | Ordinal int |
| 120 | Heading string |
| 121 | Text string |
| 122 | // Title is the document's title. The same value for every chunk of |
| 123 | // a document, as the heading trail is for every chunk of a section. |
| 124 | Title string |
| 125 | // Path is the trail of the document's ANCESTORS in its source — |
| 126 | // "Handbook › Onboarding" for a Notion page nested two deep. Empty |
| 127 | // for anything uploaded or pasted, which has no ancestry. |
| 128 | // |
| 129 | // A separate column from Heading rather than a prefix on it, and |
| 130 | // from Title rather than folded into it: PathWeight says why. |
| 131 | Path string |
| 132 | } |
| 133 | |
| 134 | // Add indexes entries in one transaction. |
| 135 | func (i *Index) Add(entries ...Entry) error { |
| 136 | i.mu.Lock() |
| 137 | defer i.mu.Unlock() |
| 138 | tx, err := i.db.Begin() |
| 139 | if err != nil { |
| 140 | return err |
| 141 | } |
| 142 | defer tx.Rollback() |
| 143 | st, err := tx.Prepare(`INSERT INTO chunks (heading, body, title, path, chunk_id, document_id, ordinal) VALUES (?, ?, ?, ?, ?, ?, ?)`) |
| 144 | if err != nil { |
| 145 | return err |
| 146 | } |
| 147 | defer st.Close() |
| 148 | for _, e := range entries { |
| 149 | if _, err := st.Exec(e.Heading, e.Text, e.Title, e.Path, e.ChunkID, e.DocumentID, e.Ordinal); err != nil { |
| 150 | return err |
| 151 | } |
| 152 | } |
| 153 | return tx.Commit() |
| 154 | } |
| 155 | |
| 156 | // RemoveDocument drops every chunk of a document. Called on delete, and |
| 157 | // before re-indexing a changed document so a stale passage cannot |
| 158 | // survive as a citation to text that is no longer there. |
| 159 | func (i *Index) RemoveDocument(documentID int64) error { |
| 160 | i.mu.Lock() |
| 161 | defer i.mu.Unlock() |
| 162 | _, err := i.db.Exec(`DELETE FROM chunks WHERE document_id = ?`, documentID) |
| 163 | return err |
| 164 | } |
| 165 | |
| 166 | // Count returns how many chunks are indexed. Used by the UI to say what |
| 167 | // the corpus actually contains rather than implying it is complete. |
| 168 | func (i *Index) Count() (int, error) { |
| 169 | i.mu.RLock() |
| 170 | defer i.mu.RUnlock() |
| 171 | var n int |
| 172 | err := i.db.QueryRow(`SELECT count(*) FROM chunks`).Scan(&n) |
| 173 | return n, err |
| 174 | } |
| 175 | |
| 176 | // Column weights for bm25(), in column order: heading, body, title, path. |
| 177 | // |
| 178 | // HeadingWeight above BodyWeight because a passage under "Revenue" is a |
| 179 | // better answer to "revenue" than one that mentions the word in passing. |
| 180 | // |
| 181 | // TitleWeight equals BodyWeight, not more. The bug it fixes is that a |
| 182 | // title term found nothing at all; ranking a title match above a |
| 183 | // document that genuinely discusses the term would be a different and |
| 184 | // worse answer. Note the consequence, which the repeated heading trail |
| 185 | // already has: a title match hits every chunk of that document, so a |
| 186 | // one-word title query can fill the page with one document. Acceptable |
| 187 | // while documents are small; revisit with §12. |
| 188 | // |
| 189 | // PathWeight is BELOW all three, and that is the load-bearing one. A |
| 190 | // document's ancestry repeats into every one of its chunks — every |
| 191 | // passage of every page under "Handbook" carries "Handbook" — so at the |
| 192 | // heading's weight the top of a source's tree outranks the content of |
| 193 | // it, and a search for "handbook" returns the whole handbook above the |
| 194 | // page that answers the question. The failure would look like bad |
| 195 | // search rather than a bug, which is the kind that survives. |
| 196 | // |
| 197 | // 0.5 IS A GUESS. It is half the body's weight because nothing has |
| 198 | // measured it; the corpus it wants tuning against does not exist yet. |
| 199 | // What is not a guess is the ORDERING, which |
| 200 | // TestAncestryCannotOutrankContent pins by measurement — and that |
| 201 | // measurement was re-taken once title joined the table, because a title |
| 202 | // column at BodyWeight is exactly the ingredient that could reintroduce |
| 203 | // the trap. See that test for the observed scores. |
| 204 | const ( |
| 205 | HeadingWeight = 3.0 |
| 206 | BodyWeight = 1.0 |
| 207 | TitleWeight = 1.0 |
| 208 | PathWeight = 0.5 |
| 209 | ) |
| 210 | |
| 211 | // Search runs a BM25 query, weighted by the constants above. |
| 212 | func (i *Index) Search(query string, limit int) ([]Hit, error) { |
| 213 | q := MatchQuery(query) |
| 214 | if q == "" { |
| 215 | return nil, nil |
| 216 | } |
| 217 | if limit <= 0 { |
| 218 | limit = 20 |
| 219 | } |
| 220 | i.mu.RLock() |
| 221 | defer i.mu.RUnlock() |
| 222 | // snippet()'s second argument is a COLUMN INDEX, and 1 is body. |
| 223 | // The path column is appended third for exactly that reason: adding |
| 224 | // it before body would silently start snippeting the wrong column, |
| 225 | // and a snippet of the wrong column still renders. |
| 226 | rows, err := i.db.Query(` |
| 227 | SELECT chunk_id, document_id, ordinal, heading, |
| 228 | snippet(chunks, 1, '<mark>', '</mark>', '…', 24), |
| 229 | bm25(chunks, ?, ?, ?, ?) |
| 230 | FROM chunks |
| 231 | WHERE chunks MATCH ? |
| 232 | ORDER BY bm25(chunks, ?, ?, ?, ?) |
| 233 | LIMIT ?`, |
| 234 | HeadingWeight, BodyWeight, TitleWeight, PathWeight, |
| 235 | q, |
| 236 | HeadingWeight, BodyWeight, TitleWeight, PathWeight, |
| 237 | limit) |
| 238 | if err != nil { |
| 239 | return nil, err |
| 240 | } |
| 241 | defer rows.Close() |
| 242 | var out []Hit |
| 243 | for rows.Next() { |
| 244 | var h Hit |
| 245 | if err := rows.Scan(&h.ChunkID, &h.DocumentID, &h.Ordinal, &h.Heading, &h.Snippet, &h.Score); err != nil { |
| 246 | return nil, err |
| 247 | } |
| 248 | h.Score = -h.Score |
| 249 | out = append(out, h) |
| 250 | } |
| 251 | return out, rows.Err() |
| 252 | } |
| 253 | |
| 254 | // MatchQuery turns what a person typed into an FTS5 MATCH expression. |
| 255 | // |
| 256 | // This exists because FTS5's query syntax is a language, and a person |
| 257 | // typing an apostrophe, a colon or the bare word "NEAR" into a search box |
| 258 | // is not writing one. Feeding raw input to MATCH produces a syntax error |
| 259 | // on ordinary English — "what's June's revenue" — which surfaces as an |
| 260 | // empty result page and reads exactly like "nothing matched". |
| 261 | // |
| 262 | // Every run of letters and digits becomes one quoted term. Quoting is |
| 263 | // what makes punctuation inert. |
| 264 | // |
| 265 | // Terms are ORed, not ANDed, and that is the load-bearing decision here. |
| 266 | // memoria is asked questions, not given keywords: "what is average |
| 267 | // revenue in June" carries four words the corpus will not contain |
| 268 | // alongside the two it will. ANDing them finds nothing, and finding |
| 269 | // nothing is indistinguishable from having nothing — the failure shape |
| 270 | // this codebase calls the worst available. OR keeps recall, and BM25 |
| 271 | // does the discriminating: a passage matching "revenue" and "June" |
| 272 | // outranks one matching only "what", because IDF makes a common word |
| 273 | // worth almost nothing. Precision comes from ranking, not from |
| 274 | // exclusion. |
| 275 | func MatchQuery(raw string) string { |
| 276 | var terms []string |
| 277 | var cur strings.Builder |
| 278 | flush := func() { |
| 279 | if cur.Len() == 0 { |
| 280 | return |
| 281 | } |
| 282 | terms = append(terms, cur.String()) |
| 283 | cur.Reset() |
| 284 | } |
| 285 | for _, r := range raw { |
| 286 | switch { |
| 287 | case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9': |
| 288 | cur.WriteRune(r) |
| 289 | case r > 127: |
| 290 | // Keep non-ASCII letters: "café", "Müller", CJK. unicode61 |
| 291 | // tokenises them; dropping them would silently narrow a |
| 292 | // search in any language but English. |
| 293 | cur.WriteRune(r) |
| 294 | default: |
| 295 | flush() |
| 296 | } |
| 297 | } |
| 298 | flush() |
| 299 | if len(terms) == 0 { |
| 300 | return "" |
| 301 | } |
| 302 | // A search box is not a document. Beyond this the OR chain costs |
| 303 | // more than the answer is worth: measured at 2.8ms for 100 terms, |
| 304 | // 97ms for 1000 and 2.25s for 5000 — and Search holds the index's |
| 305 | // read lock throughout, so one pasted essay stalls ingest for |
| 306 | // everyone. 64 is longer than any query a person types and cheap |
| 307 | // (~3ms). Terms past it are dropped, which narrows a query nobody |
| 308 | // typed rather than refusing one somebody did. |
| 309 | // |
| 310 | // §9.1 also specifies q capped at 1 KiB, which would bound this at |
| 311 | // the request instead and is not implemented; this cap does not |
| 312 | // depend on that one arriving. |
| 313 | const maxTerms = 64 |
| 314 | if len(terms) > maxTerms { |
| 315 | terms = terms[:maxTerms] |
| 316 | } |
| 317 | quoted := make([]string, 0, len(terms)+1) |
| 318 | for _, t := range terms { |
| 319 | quoted = append(quoted, `"`+t+`"`) |
| 320 | } |
| 321 | // The query's words required to appear TOGETHER, added as a further |
| 322 | // OR operand rather than replacing anything. |
| 323 | // |
| 324 | // This is the ranking half of the OR decision, and it was missing. |
| 325 | // OR keeps recall, and the comment on TestOrRecallStillRanksBestMatchFirst |
| 326 | // says the discriminating is then done by ranking — but BM25 alone |
| 327 | // cannot do it when one of the words is common. FTS5 clamps IDF to |
| 328 | // ~1e-6 for a term in most documents, so "number" contributed |
| 329 | // nothing and "phone number" ranked exactly as "phone" did. A |
| 330 | // passage with "Phone" in a short 3×-weighted heading beat the |
| 331 | // passage that actually stated the office phone number. |
| 332 | // |
| 333 | // FTS5 scores each phrase in the expression and sums them, so a |
| 334 | // passage carrying the words together earns a further contribution |
| 335 | // and rises; one carrying them scattered is untouched. Measured on |
| 336 | // the fixture in rank_test.go: the right passage 1.8670 → 4.0322, |
| 337 | // the wrong one unchanged at 2.7670, all fifteen still returned. |
| 338 | // |
| 339 | // An OR operand and not AND, and not on its own: both of those |
| 340 | // return only passages carrying the words together — one result |
| 341 | // where there were twenty-five. Narrowing recall to fix an ordering |
| 342 | // is the plausible-empty-list failure with a better score. |
| 343 | // |
| 344 | // NEAR rather than a plain phrase, and that was a review finding |
| 345 | // rather than a first instinct. A phrase is ORDERED, so "revenue |
| 346 | // june" scored a coincidental adjacency across a full stop above the |
| 347 | // passage whose heading IS Revenue — the same failure this whole |
| 348 | // function exists to fix, reintroduced by the fix. Before the phrase |
| 349 | // existed the two spellings ranked identically, because order could |
| 350 | // not matter; NEAR restores that and is what a person means by |
| 351 | // typing two words. |
| 352 | // |
| 353 | // Distance 2, measured rather than picked. On "office phone number": |
| 354 | // NEAR/0 ranked the WRONG passage first (worse than the phrase), |
| 355 | // NEAR/2 got it right, NEAR/5 promoted a loose match over a tight |
| 356 | // one. Two tolerates the small words English puts between the words |
| 357 | // someone searched for. |
| 358 | // |
| 359 | // What it cannot see is punctuation. FTS5 tokenises it away, so |
| 360 | // "…helped revenue. June also…" is adjacent here and there is no |
| 361 | // expression that says otherwise — the limit is the tokenizer's, not |
| 362 | // this line's, and the next person should not go looking at the |
| 363 | // weights for it. |
| 364 | if len(terms) > 1 { |
| 365 | quoted = append(quoted, `NEAR(`+strings.Join(quoted, " ")+`, 2)`) |
| 366 | } |
| 367 | return strings.Join(quoted, " OR ") |
| 368 | } |
| 369 | |
| 370 | // randomName gives each in-memory index its own shared-cache key. Random |
| 371 | // rather than a counter: a counter resets when the process restarts, and |
| 372 | // two indexes alive across a restart boundary in a test binary would |
| 373 | // collide again. |
| 374 | func randomName() (string, error) { |
| 375 | b := make([]byte, 8) |
| 376 | if _, err := rand.Read(b); err != nil { |
| 377 | return "", fmt.Errorf("index: name: %w", err) |
| 378 | } |
| 379 | return "memoria-index-" + hex.EncodeToString(b), nil |
| 380 | } |
| 381 | |
| 382 | // File reports the path SQLite has backing this index, empty for an |
| 383 | // in-memory database. |
| 384 | // |
| 385 | // It exists so a test can prove §3.1's ruling rather than trust the DSN |
| 386 | // string. A review changed the DSN to a file path and the entire suite |
| 387 | // stayed green with plaintext FTS5 postings on disk, because the only |
| 388 | // plaintext gate scanned the app database's directory and nothing looked |
| 389 | // at the index at all. |
| 390 | func (i *Index) File() (string, error) { |
| 391 | i.mu.RLock() |
| 392 | defer i.mu.RUnlock() |
| 393 | rows, err := i.db.Query(`PRAGMA database_list`) |
| 394 | if err != nil { |
| 395 | return "", err |
| 396 | } |
| 397 | defer rows.Close() |
| 398 | for rows.Next() { |
| 399 | var seq int |
| 400 | var name, file string |
| 401 | if err := rows.Scan(&seq, &name, &file); err != nil { |
| 402 | return "", err |
| 403 | } |
| 404 | if name == "main" { |
| 405 | return file, nil |
| 406 | } |
| 407 | } |
| 408 | return "", rows.Err() |
| 409 | } |
| 410 | |