From 64d08a358d44938e194fde4ca364b90f13f503b6 Mon Sep 17 00:00:00 2001 From: Kyle Isom Date: Sat, 26 Sep 2026 19:02:51 -0700 Subject: [PATCH] P5b review: background embedder re-queues a failed document and waits for a busy gate; no empty vectors cached Owner review of p5b (tasks 01-03 by Ornith): - background.go: `seen` was never cleared, so a document whose background embedding failed (Dixie hiccup) or was abandoned could never be queued again: every later question on it stayed BM25 for the life of the process. Dedupe now covers only queued and in-progress documents. - background.go: acquireGate gave up after one retry and dropped the document; the task said to keep retrying once a second. It now waits (background work has no deadline). - retrieval.go/background.go: a vector the server left out is not cached as an empty vector. - bm25.go: Scores on a nil *BM25 dereferenced it; returns nil. Regression tests in review_test.go; both fail on the old code. End-to-end against Dixie (local serve): Karpathy's thesis, 95 pages / 51 chunks: bm25 in 19 s, hybrid (51 embedded) after the background; RP2040 datasheet, 642 pages / 284 chunks (a 422 before P5b): bm25 in 13.5 s, all 284 embedded within 90 s, hybrid answer correct with its quote. Co-Authored-By: Claude Opus 5.5 --- internal/digest/background.go | 32 ++++++------ internal/digest/bm25.go | 5 +- internal/digest/retrieval.go | 4 +- internal/digest/review_test.go | 96 ++++++++++++++++++++++++++++++++++ 4 files changed, 118 insertions(+), 19 deletions(-) create mode 100644 internal/digest/review_test.go diff --git a/internal/digest/background.go b/internal/digest/background.go index 2447244..1cc4819 100644 --- a/internal/digest/background.go +++ b/internal/digest/background.go @@ -85,7 +85,10 @@ func (w *bgWorker) run() { w.process(job) + // Forget the document whether it finished or not: dedupe covers only queued and + // in-progress work, so a later question re-queues a document whose embedding failed. w.mu.Lock() + delete(w.seen, job.contentHash+"|"+job.callerClass) w.inProg = false w.mu.Unlock() } @@ -121,12 +124,7 @@ func (w *bgWorker) process(job embedJob) { for k := range batch { batch[k] = job.chunks[idx[b+k]] } - if !w.acquireGate() { - if e.Log != nil { - e.Log.Warn("background embed gate") - } - return - } + w.acquireGate() ectx, cancel := context.WithTimeout(context.Background(), e.timeout()) vecs, err := e.embed(ectx, batch) cancel() @@ -138,22 +136,22 @@ func (w *bgWorker) process(job embedJob) { return } for k := range batch { - e.Cache.PutEmbed(key(idx[b+k]), vecs[k]) + if k < len(vecs) && len(vecs[k]) > 0 { + e.Cache.PutEmbed(key(idx[b+k]), vecs[k]) + } } } } -// acquireGate takes the Dixie gate, retrying once per second on a refusal so interactive digests -// get turns. It returns false only if the gate cannot be taken after the first retry, so the -// document is not stuck behind a permanently busy gate. -func (w *bgWorker) acquireGate() bool { - ok, _ := w.engine.Gate.Acquire(context.Background()) - if ok { - return true +// acquireGate takes the Dixie gate, sleeping a second after each refusal (queue full or wait +// expired) so interactive digests get their turns. Background work has no deadline: it waits. +func (w *bgWorker) acquireGate() { + for { + if ok, _ := w.engine.Gate.Acquire(context.Background()); ok { + return + } + time.Sleep(time.Second) } - time.Sleep(1 * time.Second) - ok, _ = w.engine.Gate.Acquire(context.Background()) - return ok } // WaitBackground returns when the background worker is idle: the queue empty and nothing embedding. diff --git a/internal/digest/bm25.go b/internal/digest/bm25.go index 610427e..5f82e53 100644 --- a/internal/digest/bm25.go +++ b/internal/digest/bm25.go @@ -54,7 +54,10 @@ func NewBM25(docs []string) *BM25 { // Scores returns one score per document, in order, each >= 0 and never NaN or Inf. An empty // query, an empty index or a document without any query term scores 0. func (b *BM25) Scores(query string) []float64 { - if b == nil || b.nDocs == 0 { + if b == nil { + return nil + } + if b.nDocs == 0 { return make([]float64, len(b.toks)) } out := make([]float64, len(b.toks)) diff --git a/internal/digest/retrieval.go b/internal/digest/retrieval.go index 37f5ae1..01d0558 100644 --- a/internal/digest/retrieval.go +++ b/internal/digest/retrieval.go @@ -51,7 +51,9 @@ func (e *Engine) hybridOrBM25(ctx context.Context, chunks []string, question, co return joinBudget(chunks, SmallDoc), min(len(chunks), countBudget(chunks, SmallDoc)), len(chunks), "hybrid", len(cached), err } for j, i := range uncached { - e.Cache.PutEmbed(key(i), vecs[j+1]) + if len(vecs[j+1]) > 0 { // a vector the server left out is not cached as empty + e.Cache.PutEmbed(key(i), vecs[j+1]) + } } chunkVecs := make([][]float64, len(chunks)) for k, i := range cached { diff --git a/internal/digest/review_test.go b/internal/digest/review_test.go new file mode 100644 index 0000000..ca9a918 --- /dev/null +++ b/internal/digest/review_test.go @@ -0,0 +1,96 @@ +package digest + +// P5b owner review: regression tests for the background embedder. + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" +) + +// flakyDixie is fakeDixie whose /embed answers 500 to its first failN requests. +func flakyDixie(t *testing.T, failN int64) (*fakeDixie, *atomic.Int64) { + f := &fakeDixie{} + var calls atomic.Int64 + var disc atomic.Bool + mux := http.NewServeMux() + mux.HandleFunc("/embed", func(w http.ResponseWriter, r *http.Request) { + if calls.Add(1) <= failN { + http.Error(w, "boom", http.StatusInternalServerError) + return + } + var body struct { + Input []string `json:"input"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + var data []map[string]any + for i, s := range body.Input { + data = append(data, map[string]any{"index": i, "embedding": bow(s)}) + } + f.embedded.Add(int64(len(body.Input))) + _ = json.NewEncoder(w).Encode(map[string]any{"data": data}) + }) + mux.Handle("/chat", sseHandler(t, "ok", &f.lastChat, &disc)) + f.srv = httptest.NewServer(mux) + t.Cleanup(f.srv.Close) + return f, &calls +} + +// A document whose background embedding failed is queued again by a later question; it must +// not stay "seen" (and so BM25-only) for the life of the process. +func TestBackgroundRequeuesAfterFailure(t *testing.T) { + f, _ := flakyDixie(t, 1) + e := f.engine(t, NewGate(2, 8, 5*time.Second), true) + doc := longDoc(300, 250) + n := len(SplitPages(doc, ChunkTokens)) + wait := func() { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := e.WaitBackground(ctx); err != nil { + t.Fatal(err) + } + } + q := "who signed the zanzibarite manifest?" + if r := e.Run(context.Background(), doc, q, "https://ex.test/f.pdf", "fleet"); r.Retrieval != "bm25" { + t.Fatalf("first: %+v", r) + } + wait() // the first batch failed: the document was dropped + if r := e.Run(context.Background(), doc, q, "https://ex.test/f.pdf", "fleet"); r.Retrieval != "bm25" { + t.Fatalf("second: %+v, want bm25 (nothing embedded yet)", r) + } + wait() // re-queued by the second question, and embedded this time + r := e.Run(context.Background(), doc, q, "https://ex.test/f.pdf", "fleet") + if r.Retrieval != "hybrid" || r.Embedded != n { + t.Fatalf("third: %+v, want hybrid with all %d embedded", r, n) + } +} + +// The background worker waits for a busy gate; it does not give up on the document. +func TestBackgroundWaitsForBusyGate(t *testing.T) { + f := newFakeDixie(t, false) + gate := NewGate(1, 0, 50*time.Millisecond) // no queue: a busy gate refuses at once + e := f.engine(t, gate, true) + doc := longDoc(100, 0) + n := len(SplitPages(doc, ChunkTokens)) + if ok, _ := gate.Acquire(context.Background()); !ok { + t.Fatal("could not take the gate") + } + // Run takes the gate before it chooses a retrieval mode, so queue the document directly. + chunks := SplitPages(doc, ChunkTokens) + e.queueForEmbed(chunks, "hash-g", "fleet") + time.Sleep(2500 * time.Millisecond) // longer than the old single retry + gate.Release() + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := e.WaitBackground(ctx); err != nil { + t.Fatal(err) + } + if got := f.embedded.Load(); got != int64(n) { + t.Fatalf("embedded %d of %d chunks: the worker gave up on a busy gate", got, n) + } +}