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:
@@ -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.
|
||||
Reference in New Issue
Block a user