oficina / memoria Public

Clone
git clone https://amadan.net/oficina/memoria

Plain git — no account needed to clone.

Download

Download this file

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.
13package index
14
15import (
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.
28type 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.
41type 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.
48func 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.
108func (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.
116type 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.
135func (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.
159func (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.
168func (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.
204const (
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.
212func (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.
275func 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.
374func 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.
390func (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