kyleandClaude Opus 5.5
64d08a358d
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 <noreply@anthropic.com>
97 lines
3.2 KiB
Go
97 lines
3.2 KiB
Go
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)
|
|
}
|
|
}
|