sift/docs/plans/p5b/02-retrieval.md

P5b task 02: cached embeddings and BM25-first retrieval for long documents

Branch: p5b (git switch p5b; git status --short must be empty, otherwise stop) Commit subject: digest: cache chunk embeddings; long documents answer from BM25 while embedding in the background

Goal

Today selectText embeds every chunk of a document for every question, inline, inside the digest's 45 s timeout, and never uses the embedding cache. After this task:

• chunk embeddings are cached;
• a document with few uncached chunks is embedded inline and scored hybrid (as today);
• a long one is answered at once from BM25, with no embedding call on the request path, while its chunks are embedded by a background worker;
• once they are, the next question is hybrid.

Files

• Copy: internal/digest/retrieval_test.go
• Create: internal/digest/retrieval.go (selection), internal/digest/background.go (the worker)
• Modify: internal/digest/digest.go (Result fields, Run, selectText's call site), docs/implementer-log.md

Interfaces

package digest

const (
	EmbedSyncMax = 48   // uncached chunks a question may embed inline; more go to the background
	MaxDocChunks = 5000 // background embedding cap per document
)

// in Result:
	Retrieval string // "whole", "head", "hybrid", "bm25" or "cached" (see rule 2)
	Embedded  int    // chunks that had an embedding when the answer was built

// WaitBackground returns when the background worker is idle with nothing queued, or ctx.Err().
func (e *Engine) WaitBackground(ctx context.Context) error

What already exists and stays:

• Split, SplitPages, NewBM25 (task 01);
• EstTokens, SmallDoc, ChunkTokens, EmbedBatch, lexicalOverlap, cosine;
• e.embed(ctx, inputs) (batches of EmbedBatch, halves on "too large");
• e.Gate.Acquire/Release;
• dcache.EmbedKey(contentHash string, chunkIndex int, embedModel, callerClass string), e.Cache.GetEmbed(key) / PutEmbed(key, vec) (a nil *DB is a no-op).

Rules

1. Chunks and keys.
   • Chunk with SplitPages(cleaned, ChunkTokens).
   • Embedding keys are dcache.EmbedKey(contentHash+"/sp1200", i, e.embedModel(), callerClass): the suffix binds the key to this chunker, so a future chunking change cannot reuse wrong vectors.
   • contentHash and callerClass come from Run (pass them to the selection).
2. Selection (retrieval.go), in order:
   • EstTokens(cleaned) <= SmallDoc → the whole text, Retrieval "whole".
   • Empty question → the leading chunks within SmallDoc, as today, Retrieval "head".
   • No cache (e.Cache == nil):
     • len(chunks) <= EmbedSyncMax → embed everything inline, as today, "hybrid";
     • otherwise "bm25", with no embedding and nothing queued (there is nowhere to keep the vectors).
   • Cache: look up every chunk's vector.
     • uncached <= EmbedSyncMax → embed the question and the uncached chunks inline (the question first, then the missing chunks, through e.embed), store each new chunk vector, "hybrid".
     • otherwise → "bm25": rank by NewBM25(chunks).Scores(question) normalised by its maximum. Make no embedding call, and queue the document for the background (rule 4).
   • Hybrid score = 0.7·cosine(q, v_i) + 0.3·lexicalOverlap(question, chunk), as today.
   • Both modes then use today's greedy budget: chunk 0 first, then best-scoring chunks while the total stays within SmallDoc, emitted in document order.
   • Embedded = the chunks with a vector at that moment.
3. Run:
   • Set r.Retrieval and r.Embedded.
   • On a digest-cache hit: Retrieval "cached".
   • A digest made in "bm25" mode is not stored in the digest cache. It is provisional; once the embeddings exist the same question must get a hybrid answer.
4. Background worker (background.go):
   • One goroutine per Engine, started on first use (sync.Once), with a queue of at most 4 documents, deduplicated by contentHash+"|"+callerClass. A full queue drops the document, logs Warn (with no text) and a later question re-queues it.
   • For each document: embed the first MaxDocChunks chunks that lack a vector, EmbedBatch at a time. For each batch:
     • e.Gate.Acquire(context.Background()), so interactive digests get turns; on a refusal (queue full or wait), sleep 1 s and try again;
     • e.embed with a e.timeout() context;
     • Release;
     • store the vectors.
   • On an embedding error: log Warn, drop the rest of that document, go on to the next.
   • WaitBackground is true-idle: the queue is empty and no document is in progress. Poll every 20 ms or use a condition variable; either is fine.
5. Never log document text. The worker never holds the gate between batches.

Facts

• The given test's fake dixie counts every embedded input. It can hold /embed shut, which is how it proves the "bm25" answer made no embedding call on the request path: if yours does, the first Run blocks and the test fails on time.
• sseHandler and chatReq come from digest_test.go / digest.go in the same package.
• Existing tests (TestChunkOrderAndBudget, TestEmbedTooLargeRetriesAndCompletes, …) run with no cache and must stay green.

Steps

• 1. git switch p5b; copy the given test. 2. See it fail.
• 3. retrieval.go, 4. background.go, 5. the Run/Result changes.
• 6. gofmt -w; go test -race -count=3 ./internal/digest/ → ok. 7. make gate.
• 8. Row p5b/02-retrieval; commit.

git add internal/digest docs/implementer-log.md
git commit -F .state/commit-msg.txt

Done when

• The given and existing digest tests pass under -race -count=3; make gate ok; given file byte-identical.

Stop and report if

• A given test's timing cannot be met without sleeping on the request path.