P5b and P5c plans: long-document retrieval (BM25 first, background embedding) and /v1/extract + sift extract; given tests

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
2026-09-26 11:01:18 -07:00
co-authored by Claude Opus 5.5
parent 1311cf38d4
commit 49d96e197b
14 changed files with 1464 additions and 0 deletions
+105
View File
@@ -0,0 +1,105 @@
# 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
```go
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.
```sh
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.