diff --git a/docs/implementer-log.md b/docs/implementer-log.md index ee49493..db8dfed 100644 --- a/docs/implementer-log.md +++ b/docs/implementer-log.md @@ -12,6 +12,7 @@ One row per task session that ends in a commit. Newest last. | p5a/06-api-pdf | 2026-09-26 | done | resumed | gate: ok (7 of 10 PDF tests; 3 are task 07's) | none | Served PDFs through `/v1/fetch` (clean) and `/v2/scrape`. Added `RetryAfter int` to `fetch.Error` and a `Retry-After` header line guarded by `> 0`. Added `PDFRunner`/`PDFPool`/`PDFTTL`/`PDFMaxQuestionChunks` Server fields (PDFPool built once via `sync.Once` only when nil, so the test's direct override is respected) and `pdfPoolOnce`. Routed `clean.Kind` in `doFetch`/`cleanFetched` (html/json/text switched, "other"→415, "pdf"→extractPDF); `doFetch`'s switch routes "text" explicitly so a PDF never runs `clean.Text` first. New `internal/api/pdf.go`: `DefaultPDFTTL`/`DefaultPDFMaxQuestionChunks` (task 07), `runner()`/`pool()` defaults, `extractPDF` (pool key `callerClass|SHA256`, two-value assert, `fmt.Errorf("pdf: bad result %T")`), `pdfError` (422 for the five pdftext errors + panic default, 503+RetryAfter 10 for ErrBusy, 429 for ErrCallerBusy), `renderPDF` (`pdfdoc.Render`, whole text or map). `doFetch` returns the 4th `doc *pdftext.Doc`; digest uses `pdfdoc.Full(doc)` (computed once) as content, clean uses `renderPDF.Content`, envelope `pdf` object has pages/title/outline/extracted/first/last/map/grep_hits/note, raw mode→400. Scrape markdown=`renderPDF.Content`, html="", `metadata.numPages`=doc.Pages. 7 of 10 given tests pass under -race -count=3; the 3 failures (`TestPDFPagesAndGrep`, `TestPDFCache`, `TestPDFQuestionTooLong`) are task 07 tests that reference fields not in this task, exactly as the task doc (lines 145, 157-158) predicts; nothing else fails. Given files byte-identical. | | p5a/08-digest-cli-docs | 2026-09-26 | done | 1 | gate: ok | none | Added the page-cite sentence to the digest `systemPrompt` ("When the document has `--- page N ---` markers, cite the page of each fact as (p. N).") and set `PromptVersion = "3"`. Added the `sift_pdf_extract_total{outcome}`, `sift_pdf_extract_seconds`, `sift_pdf_pages_total`, `sift_pdf_pool_wait_seconds`, `sift_pdf_cache_hits_total` series to `internal/metrics` (`pdf.go` package) and wired `ObservePDF`/`ObservePDFWait`/`PDFCacheHit` into `internal/api/pdf.go` (`ObservePDF` once per call after `extractPDF`, outcome classified by `pdfOutcome` so raw text never becomes a label; `ObservePDFWait` only when the caller actually waited on `pool().Do`; `PDFCacheHit` on fresh hits and 304s). Refactored `cmd/sift/fetchcmd.go` into a `fetchOpts` struct with `parseFetchOpts` (returns `fetchOpts`, not `string, error`) and a `fetchBody(o fetchOpts) map[string]any` that does the digest/clean default and adds `pages`/`grep` only when set; `--pages`/`--grep` added to `fetchValueFlags`. Updated `skill/SKILL.md` (PDFs section; removed PDF mention from the 415 row). Unit: `nix/module.nix` `MemoryMax = "1G"` (one line). README: added a "PDF documents" section covering detection, limits (100 MB/120 s download, 60 s extraction, 2,000 pages), the child process, the 2-slot/4-wait/20 s pool, the 24 h conditional-GET cache, the answers, the task-06 errors table, and the metrics; noted the `go-pdfkit/extract` v0.3.0 engine chosen by the `docs/P5a.md` §3 bake-off. All tests pass under -race, make gate ok, three given files byte-identical, nix diff is one line. | | p5b/01-splitpages-bm25 | 2026-09-26 | done | 1 | gate: ok | none | `splitpages.go`: `SplitPages` is Split's budget/line algorithm plus page awareness — a `--- page N ---` line sets the current page, and a chunk that starts after the first marker but not on a marker line is prefixed `--- page N (cont.) ---\n\n` for the page its first line belongs to (captured at chunk start, so a chunk that crosses a marker is still labelled with its starting page; the marker counts toward the chunk budget). No-marker text yields exactly Split's result, so `TestSplitKeepsPageMarkers` is untouched. `bm25.go`: `BM25` with k1=1.2/b=0.75, term counts via a local `docTerms` (same tokenization as `terms()`, >=2 chars), IDF `ln(1+(N-df+0.5)/(df+0.5))`, avgDL over all docs (guarded to 1 when the corpus is all-empty), one score per doc >= 0 with no query term scoring 0 and empty query/index scoring 0 everywhere. All digest tests pass under -race -count=3, make gate ok, two given files byte-identical. | +| p5b/02-retrieval | 2026-09-26 | done | 1 | gate: ok | none | Cached chunk embeddings + BM25-first retrieval. `digest.go`: `EmbedSyncMax=48`, `MaxDocChunks=5000`; `Result` gains `Retrieval string` and `Embedded int`; `selectText` now returns `(body, used, total, mode, embedded, err)` and `Run` sets both; `PutDigest` skips storing any `mode=="bm25"` answer (provisional, not cached). `retrieval.go`: `hybridOrBM25` — no cache, `len(chunks)<=48` embeds question+chunks inline and scores hybrid; `>48` scores BM25 and queues; cache present embeds the question and only the uncached chunks inline and scores hybrid on the union; both embed paths return hybrid, both bm25 paths return `bm25`. `background.go`: `bgWorker` with `bgOnce sync.Once` (creates the worker, never starts a goroutine) + `w.started` under `w.mu` (starts the goroutine once), dedup by `contentHash|callerClass`, full-queue drops the doc without marking seen so a later question re-queues, `acquireGate` retries once after 1 s on refusal, `WaitBackground` polls every 20 ms for true idle. `retrieval_test.go` copied byte-identical from the plan. All 12 digest tests pass under `-race -count=3`; `make gate` ok. | | p5a/07-api-cache-pages | 2026-09-26 | done | 1 | gate: ok | none | Implemented the PDF document cache with revalidation in `doFetch`: `pdfFromCache` (miss / fresh within `PDFTTL` / stale → `FetchIf` with 304→`TouchPDF`+answer-fresh / 200→return response with ok=false and no second GET) and `storePDF` after `extractPDF` (`ContentHash=res.SHA256`, ETag/Last-Modified from headers, FetchedAt=now, Meta=JSON of `doc.Meta`, Pages=`doc.Text`). `pages`/`grep` passed into `pdfdoc.Options` in `renderPDF`; `ErrBadPages`/`ErrBadGrep`→400 with the error's own text, non-PDF→400 "pages and grep apply to PDFs only" before the doc block. Question limit in digest mode: `n=len(digest.Split(pdfdoc.Full(doc), ChunkTokens))`, `n>PDFMaxQuestionChunks`→422 "document too long for question (n chunks); use grep or pages" before any digest call. Envelope `cached` set true for cache answers via the synthetic cache res's empty body (`len(res.Body)==0`); the one `doFetch` change beyond the rule-1 call was reusing that 200 response instead of a second `Fetch`. Added `fmt` import to api.go. Given file byte-identical; make gate ok. | ## P5a review (owner, 2026-09-26) diff --git a/internal/digest/background.go b/internal/digest/background.go new file mode 100644 index 0000000..2447244 --- /dev/null +++ b/internal/digest/background.go @@ -0,0 +1,174 @@ +package digest + +import ( + "context" + "sync" + "time" + + "git.wntrmute.dev/kyle/sift/internal/dcache" +) + +// maxBgQueue is the number of documents the background worker may hold at once. +const maxBgQueue = 4 + +// bgWorker is the single background goroutine that embeds the chunks of a long document after a +// bm25 answer queued it. It owns a bounded, deduplicated queue and yields the Dixie gate to +// interactive digests one batch at a time. +type bgWorker struct { + engine *Engine + + mu sync.Mutex + cond *sync.Cond + queue []embedJob + seen map[string]struct{} + inProg bool + started bool +} + +type embedJob struct { + chunks []string + contentHash string + callerClass string +} + +// worker returns the engine's background worker, creating it once. It does not start the goroutine. +func (e *Engine) worker() *bgWorker { + e.bgOnce.Do(func() { + w := &bgWorker{engine: e} + w.cond = sync.NewCond(&w.mu) + w.seen = map[string]struct{}{} + e.bg = w + }) + return e.bg +} + +// queueForEmbed asks the background worker to embed a document's chunks. The document is +// deduplicated by contentHash+"|"+callerClass; a full queue drops it (and logs) so a later +// question re-queues it. +func (e *Engine) queueForEmbed(chunks []string, contentHash, callerClass string) { + if e == nil { + return + } + w := e.worker() + w.mu.Lock() + defer w.mu.Unlock() + if !w.started { + w.started = true + go w.run() + } + key := contentHash + "|" + callerClass + if _, ok := w.seen[key]; ok { + return + } + if len(w.queue) >= maxBgQueue { + // Leave it un-seen: a later question re-queues it. + if e.Log != nil { + e.Log.Warn("background embed queue full") + } + return + } + w.seen[key] = struct{}{} + w.queue = append(w.queue, embedJob{chunks: chunks, contentHash: contentHash, callerClass: callerClass}) + w.cond.Signal() +} + +func (w *bgWorker) run() { + for { + w.mu.Lock() + for len(w.queue) == 0 { + w.cond.Wait() + } + job := w.queue[0] + w.queue = w.queue[1:] + w.inProg = true + w.mu.Unlock() + + w.process(job) + + w.mu.Lock() + w.inProg = false + w.mu.Unlock() + } +} + +// process embeds the first MaxDocChunks chunks that lack a vector, EmbedBatch at a time, yielding +// the gate between batches and storing each vector. On any embedding error it drops the rest of +// the document. +func (w *bgWorker) process(job embedJob) { + e := w.engine + key := func(i int) string { + return dcache.EmbedKey(job.contentHash+"/sp1200", i, e.embedModel(), job.callerClass) + } + var idx []int + for i := range job.chunks { + if i >= MaxDocChunks { + break + } + if _, ok := e.Cache.GetEmbed(key(i)); ok { + continue + } + idx = append(idx, i) + if len(idx) >= MaxDocChunks { + break + } + } + for b := 0; b < len(idx); b += EmbedBatch { + j := b + EmbedBatch + if j > len(idx) { + j = len(idx) + } + batch := make([]string, j-b) + 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 + } + ectx, cancel := context.WithTimeout(context.Background(), e.timeout()) + vecs, err := e.embed(ectx, batch) + cancel() + e.Gate.Release() + if err != nil { + if e.Log != nil { + e.Log.Warn("background embed", "err", truncateErr(err, 512)) + } + return + } + for k := range batch { + 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 + } + 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. +func (e *Engine) WaitBackground(ctx context.Context) error { + w := e.worker() + for { + w.mu.Lock() + idle := len(w.queue) == 0 && !w.inProg + w.mu.Unlock() + if idle { + return nil + } + if err := ctx.Err(); err != nil { + return err + } + time.Sleep(20 * time.Millisecond) + } +} diff --git a/internal/digest/digest.go b/internal/digest/digest.go index 1c0ca10..5df263d 100644 --- a/internal/digest/digest.go +++ b/internal/digest/digest.go @@ -33,6 +33,8 @@ const ( ChunkTokens = 1200 EmbedBatch = 16 EmbedMaxRunes = 3000 + EmbedSyncMax = 48 + MaxDocChunks = 5000 DefaultPermit = 2 DefaultQDepth = 8 DefaultQWait = 20 * time.Second @@ -191,6 +193,9 @@ type Engine struct { Log *slog.Logger Cache *dcache.DB + bgOnce sync.Once + bg *bgWorker + mu sync.Mutex Calls uint64 Fallback map[string]uint64 @@ -205,6 +210,8 @@ type Result struct { DigestMS int64 ReturnedTok int Cached bool + Retrieval string // "whole", "head", "hybrid", "bm25" or "cached" + Embedded int // chunks that had an embedding when the answer was built } func (e *Engine) timeout() time.Duration { @@ -305,6 +312,7 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL, callerCla r.ChunksUsed = h.ChunksUsed r.ReturnedTok = EstTokens(h.Text) r.Cached = true + r.Retrieval = "cached" r.DigestMS = time.Since(start).Milliseconds() return r } @@ -327,8 +335,10 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL, callerCla dctx, cancel := context.WithTimeout(ctx, e.timeout()) defer cancel() - body, nUsed, nTotal, err := e.selectText(dctx, cleaned, question) + body, nUsed, nTotal, mode, embedded, err := e.selectText(dctx, cleaned, question, contentHash, callerClass) r.ChunksUsed, r.ChunksTotal = nUsed, nTotal + r.Retrieval = mode + r.Embedded = embedded if err != nil { e.logFallback("embed", err) r.Fallback = "embed" @@ -359,72 +369,26 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL, callerCla } r.Text = text r.ReturnedTok = EstTokens(text) - if e != nil && e.Cache != nil { + // A bm25 answer is provisional: once the embeddings exist the same question + // must get a hybrid answer, so it does not go in the digest cache. + if e != nil && e.Cache != nil && mode != "bm25" { e.Cache.PutDigest(dkey, text, r.ChunksUsed) } return r } -func (e *Engine) selectText(ctx context.Context, cleaned, question string) (string, int, int, error) { +func (e *Engine) selectText(ctx context.Context, cleaned, question, contentHash, callerClass string) (string, int, int, string, int, error) { if EstTokens(cleaned) <= SmallDoc { - return cleaned, 1, 1, nil + return cleaned, 1, 1, "whole", 0, nil } - chunks := Split(cleaned, ChunkTokens) + chunks := SplitPages(cleaned, ChunkTokens) if len(chunks) == 0 { - return cleaned, 1, 1, nil + return cleaned, 1, 1, "whole", 0, nil } if question == "" { - return joinBudget(chunks, SmallDoc), min(len(chunks), countBudget(chunks, SmallDoc)), len(chunks), nil + return joinBudget(chunks, SmallDoc), min(len(chunks), countBudget(chunks, SmallDoc)), len(chunks), "head", 0, nil } - vecs, err := e.embed(ctx, append([]string{question}, capEmbed(chunks)...)) - if err != nil || len(vecs) < 2 { - return joinBudget(chunks, SmallDoc), min(len(chunks), countBudget(chunks, SmallDoc)), len(chunks), err - } - q := vecs[0] - type scored struct { - i int - s float64 - } - var ss []scored - for i, c := range chunks { - if i+1 >= len(vecs) { - break - } - ss = append(ss, scored{i: i, s: 0.7*cosine(q, vecs[i+1]) + 0.3*lexicalOverlap(question, c)}) - } - picked := map[int]struct{}{} - used := 0 - if len(chunks) > 0 { - picked[0] = struct{}{} - used = EstTokens(chunks[0]) - } - for len(picked) < len(ss) { - best, bi := -1.0, -1 - for _, x := range ss { - if _, ok := picked[x.i]; ok { - continue - } - if x.s > best { - best, bi = x.s, x.i - } - } - if bi < 0 { - break - } - need := EstTokens(chunks[bi]) - if used+need > SmallDoc && used > 0 { - break - } - picked[bi] = struct{}{} - used += need - } - var ordered []string - for i := range chunks { - if _, ok := picked[i]; ok { - ordered = append(ordered, chunks[i]) - } - } - return strings.Join(ordered, "\n\n"), len(ordered), len(chunks), nil + return e.hybridOrBM25(ctx, chunks, question, contentHash, callerClass) } func capEmbed(chunks []string) []string { diff --git a/internal/digest/digest_test.go b/internal/digest/digest_test.go index cbbcdfe..a0373ad 100644 --- a/internal/digest/digest_test.go +++ b/internal/digest/digest_test.go @@ -374,7 +374,7 @@ func TestPinChunk0Infobox(t *testing.T) { if EstTokens(cleaned) <= SmallDoc { t.Fatalf("fixture too small: %d", EstTokens(cleaned)) } - body, nUsed, nTotal, err := e.selectText(context.Background(), cleaned, "when was the initial release?") + body, nUsed, nTotal, _, _, err := e.selectText(context.Background(), cleaned, "when was the initial release?", "", dcache.ClassFleet) if err != nil { t.Fatal(err) } diff --git a/internal/digest/retrieval.go b/internal/digest/retrieval.go new file mode 100644 index 0000000..37f5ae1 --- /dev/null +++ b/internal/digest/retrieval.go @@ -0,0 +1,147 @@ +package digest + +import ( + "context" + "strings" + + "git.wntrmute.dev/kyle/sift/internal/dcache" +) + +// hybridOrBM25 scores a chunked document either hybrid (cosine + lexical overlap) or, when there +// are too many uncached chunks to embed inline, lexical-only BM25. mode is "hybrid" or "bm25". +func (e *Engine) hybridOrBM25(ctx context.Context, chunks []string, question, contentHash, callerClass string) (string, int, int, string, int, error) { + key := func(i int) string { + return dcache.EmbedKey(contentHash+"/sp1200", i, e.embedModel(), callerClass) + } + + if e.Cache == nil { + if len(chunks) <= EmbedSyncMax { + vecs, err := e.embed(ctx, append([]string{question}, capEmbed(chunks)...)) + if err != nil { + return joinBudget(chunks, SmallDoc), min(len(chunks), countBudget(chunks, SmallDoc)), len(chunks), "hybrid", 0, err + } + chunkVecs := make([][]float64, len(chunks)) + for i := range chunks { + chunkVecs[i] = vecs[i+1] + } + body, used := greedyBudget(chunks, hybridScores(chunks, question, vecs[0], chunkVecs), SmallDoc) + return body, used, len(chunks), "hybrid", len(chunks), nil + } + body, used := greedyBudget(chunks, bm25Scores(chunks, question), SmallDoc) + return body, used, len(chunks), "bm25", 0, nil + } + + // Cache present: look up every chunk's vector. + var cached []int + var cachedVecs [][]float64 + var uncached []int + for i := range chunks { + if v, ok := e.Cache.GetEmbed(key(i)); ok { + cached = append(cached, i) + cachedVecs = append(cachedVecs, v) + continue + } + uncached = append(uncached, i) + } + + if len(uncached) <= EmbedSyncMax { + in := append([]string{question}, capEmbed(uncachedChunks(chunks, uncached))...) + vecs, err := e.embed(ctx, in) + if err != nil { + 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]) + } + chunkVecs := make([][]float64, len(chunks)) + for k, i := range cached { + chunkVecs[i] = cachedVecs[k] + } + for j, i := range uncached { + chunkVecs[i] = vecs[j+1] + } + body, used := greedyBudget(chunks, hybridScores(chunks, question, vecs[0], chunkVecs), SmallDoc) + return body, used, len(chunks), "hybrid", len(chunks), nil + } + + body, used := greedyBudget(chunks, bm25Scores(chunks, question), SmallDoc) + e.queueForEmbed(chunks, contentHash, callerClass) + return body, used, len(chunks), "bm25", len(cached), nil +} + +func uncachedChunks(chunks []string, uncached []int) []string { + out := make([]string, len(uncached)) + for j, i := range uncached { + out[j] = chunks[i] + } + return out +} + +// hybridScores builds the hybrid scorer 0.7*cosine(q, v_i) + 0.3*lexicalOverlap(question, chunk), +// where chunkVecs[i] is chunk i's vector. +func hybridScores(chunks []string, question string, q []float64, chunkVecs [][]float64) func(int) float64 { + return func(i int) float64 { + return 0.7*cosine(q, chunkVecs[i]) + 0.3*lexicalOverlap(question, chunks[i]) + } +} + +// bm25Scores builds the BM25 scorer, normalised by its maximum so scores land in [0, 1]. +func bm25Scores(chunks []string, question string) func(int) float64 { + scores := NewBM25(chunks).Scores(question) + var max float64 + for _, s := range scores { + if s > max { + max = s + } + } + if max == 0 { + return func(int) float64 { return 0 } + } + return func(i int) float64 { return scores[i] / max } +} + +// greedyBudget picks chunk 0 first, then the highest-scoring remaining chunks while the total +// stays within budget tokens, and returns them in document order with the count used. +func greedyBudget(chunks []string, score func(int) float64, budget int) (string, int) { + type scored struct { + i int + s float64 + } + var ss []scored + for i := range chunks { + ss = append(ss, scored{i: i, s: score(i)}) + } + picked := map[int]struct{}{} + used := 0 + if len(chunks) > 0 { + picked[0] = struct{}{} + used = EstTokens(chunks[0]) + } + for len(picked) < len(ss) { + best, bi := -1.0, -1 + for _, x := range ss { + if _, ok := picked[x.i]; ok { + continue + } + if x.s > best { + best, bi = x.s, x.i + } + } + if bi < 0 { + break + } + need := EstTokens(chunks[bi]) + if used+need > budget && used > 0 { + break + } + picked[bi] = struct{}{} + used += need + } + var ordered []string + for i := range chunks { + if _, ok := picked[i]; ok { + ordered = append(ordered, chunks[i]) + } + } + return strings.Join(ordered, "\n\n"), len(ordered) +} diff --git a/internal/digest/retrieval_test.go b/internal/digest/retrieval_test.go new file mode 100644 index 0000000..c7193c0 --- /dev/null +++ b/internal/digest/retrieval_test.go @@ -0,0 +1,247 @@ +package digest + +// P5b task 02: retrieval for long documents. Embeddings are cached per chunk; a document with +// more uncached chunks than EmbedSyncMax is answered at once from BM25 while its chunks are +// embedded in the background, and later questions use the hybrid score. + +import ( + "context" + "encoding/json" + "fmt" + "hash/fnv" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "git.wntrmute.dev/kyle/sift/internal/dcache" +) + +// fakeDixie serves /embed (a bag-of-words hash into 64 dims, so cosine similarity means +// something) and /chat (the SSE stub). It counts embedded inputs and can hold /embed shut. +type fakeDixie struct { + srv *httptest.Server + embedded atomic.Int64 // inputs embedded, questions included + hold chan struct{} + holdOnce sync.Once + lastChat atomic.Pointer[chatReq] +} + +func newFakeDixie(t *testing.T, hold bool) *fakeDixie { + f := &fakeDixie{} + if hold { + f.hold = make(chan struct{}) + } + var disc atomic.Bool + mux := http.NewServeMux() + mux.HandleFunc("/embed", func(w http.ResponseWriter, r *http.Request) { + if f.hold != nil { + select { + case <-f.hold: + case <-r.Context().Done(): + 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(func() { f.release(); f.srv.Close() }) + return f +} + +func (f *fakeDixie) release() { + if f.hold != nil { + f.holdOnce.Do(func() { close(f.hold) }) + } +} + +func bow(s string) []float64 { + v := make([]float64, 64) + for t := range terms(s) { + h := fnv.New32a() + _, _ = h.Write([]byte(t)) + v[h.Sum32()%64]++ + } + return v +} + +func (f *fakeDixie) engine(t *testing.T, gate *Gate, withCache bool) *Engine { + e := &Engine{ChatURL: f.srv.URL + "/chat", EmbedURL: f.srv.URL + "/embed", Gate: gate, Client: f.srv.Client()} + if withCache { + db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 0, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + e.Cache = db + } + return e +} + +// longDoc is n pages of distinct filler; the page named rare carries a rare term. +func longDoc(n, rare int) string { + var b strings.Builder + words := []string{"harbour", "vessel", "cargo", "anchor", "pilot", "tug", "berth", "quay", "hull", "keel", "mast", "deck"} + for p := 1; p <= n; p++ { + if p > 1 { + b.WriteString("\n\n") + } + fmt.Fprintf(&b, "--- page %d ---\n\n", p) + for l := 0; l < 40; l++ { + fmt.Fprintf(&b, "page %d line %d %s %s %s %s\n", p, l, words[(p+l)%12], words[(p*3+l)%12], words[(p*5+l)%12], words[(p*7+l)%12]) + } + if p == rare { + b.WriteString("the zanzibarite cargo manifest was signed by the second officer\n") + } + } + return b.String() +} + +func chatDoc(f *fakeDixie) string { + req := f.lastChat.Load() + if req == nil { + return "" + } + var s strings.Builder + for _, m := range req.Messages { + s.WriteString(m.Content) + } + return s.String() +} + +func TestEmbeddingsAreCachedAcrossQuestions(t *testing.T) { + f := newFakeDixie(t, false) + e := f.engine(t, NewGate(2, 8, 5*time.Second), true) + doc := longDoc(30, 0) // ~11.5k tokens: over SmallDoc (so not "whole"), ~10 chunks: under EmbedSyncMax + if EstTokens(doc) <= SmallDoc { + t.Fatalf("test document is %d tokens; it must exceed SmallDoc (%d) to reach retrieval", EstTokens(doc), SmallDoc) + } + r := e.Run(context.Background(), doc, "which berth did the tug use?", "https://ex.test/a.pdf", dcache.ClassFleet) + if r.Fallback != "" || r.Retrieval != "hybrid" { + t.Fatalf("first question: fallback %q retrieval %q", r.Fallback, r.Retrieval) + } + n := len(SplitPages(doc, ChunkTokens)) + if got := f.embedded.Load(); got != int64(n+1) { + t.Fatalf("first question embedded %d inputs, want %d chunks + the question", got, n) + } + r = e.Run(context.Background(), doc, "where was the pilot?", "https://ex.test/a.pdf", dcache.ClassFleet) + if r.Fallback != "" || r.Retrieval != "hybrid" || r.Embedded != n { + t.Fatalf("second question: %+v", r) + } + if got := f.embedded.Load(); got != int64(n+2) { + t.Errorf("second question embedded %d inputs in total, want %d: only its own question is new", got, n+2) + } +} + +// A long document is answered from BM25 without waiting for a single embedding, and becomes +// hybrid once the background has embedded every chunk. +func TestLongDocumentBM25ThenHybrid(t *testing.T) { + f := newFakeDixie(t, true) // /embed held shut until release + e := f.engine(t, NewGate(2, 8, 5*time.Second), true) + doc := longDoc(300, 250) + n := len(SplitPages(doc, ChunkTokens)) + if n <= EmbedSyncMax { + t.Fatalf("test document has %d chunks; it must exceed EmbedSyncMax (%d)", n, EmbedSyncMax) + } + start := time.Now() + r := e.Run(context.Background(), doc, "who signed the zanzibarite manifest?", "https://ex.test/big.pdf", dcache.ClassFleet) + if d := time.Since(start); d > 10*time.Second { + t.Fatalf("first answer took %v: it waited on embeddings", d) + } + if r.Fallback != "" || r.Retrieval != "bm25" || r.Embedded != 0 { + t.Fatalf("first answer: %+v, want retrieval bm25 with 0 embedded", r) + } + if doc := chatDoc(f); !strings.Contains(doc, "zanzibarite") || !strings.Contains(doc, "page 250") { + t.Errorf("BM25 did not bring page 250 (the only one with the term) into the prompt") + } + + f.release() + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + if err := e.WaitBackground(ctx); err != nil { + t.Fatalf("background embedding did not finish: %v", err) + } + r = e.Run(context.Background(), doc, "who signed the zanzibarite manifest?", "https://ex.test/big.pdf", dcache.ClassFleet) + if r.Fallback != "" || r.Retrieval != "hybrid" || r.Embedded != n { + t.Fatalf("after background: %+v, want hybrid with all %d embedded", r, n) + } + if doc := chatDoc(f); !strings.Contains(doc, "zanzibarite") { + t.Errorf("hybrid lost the page with the term") + } +} + +// Two questions on the same long document queue it once: every chunk is embedded once. +func TestBackgroundDedupes(t *testing.T) { + f := newFakeDixie(t, true) + e := f.engine(t, NewGate(2, 8, 5*time.Second), true) + doc := longDoc(200, 0) + n := len(SplitPages(doc, ChunkTokens)) + e.Run(context.Background(), doc, "first question about the quay", "https://ex.test/b.pdf", dcache.ClassFleet) + e.Run(context.Background(), doc, "second question about the keel", "https://ex.test/b.pdf", dcache.ClassFleet) + f.release() + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + if err := e.WaitBackground(ctx); err != nil { + t.Fatal(err) + } + if got := f.embedded.Load(); got > int64(n) { + t.Errorf("embedded %d inputs for a %d-chunk document: queued twice", got, n) + } +} + +// Background work takes the gate one batch at a time, so an interactive digest is not starved. +func TestBackgroundYieldsTheGate(t *testing.T) { + f := newFakeDixie(t, false) + e := f.engine(t, NewGate(1, 8, 10*time.Second), true) + e.Run(context.Background(), longDoc(400, 0), "start the background", "https://ex.test/c.pdf", dcache.ClassFleet) + start := time.Now() + r := e.Run(context.Background(), "a short page about a tug and a quay", "what is here?", "https://ex.test/small", dcache.ClassFleet) + if r.Fallback != "" { + t.Fatalf("interactive digest fell back (%q) while background embedding ran", r.Fallback) + } + if d := time.Since(start); d > 5*time.Second { + t.Errorf("interactive digest waited %v behind background embedding", d) + } + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + _ = e.WaitBackground(ctx) +} + +// Without a cache there is nowhere to keep background work: a long document is BM25 and +// nothing is embedded. +func TestNoCacheLongDocumentIsBM25Only(t *testing.T) { + f := newFakeDixie(t, false) + e := f.engine(t, NewGate(2, 8, 5*time.Second), false) + r := e.Run(context.Background(), longDoc(300, 250), "who signed the zanzibarite manifest?", "https://ex.test/d.pdf", dcache.ClassFleet) + if r.Fallback != "" || r.Retrieval != "bm25" { + t.Fatalf("%+v, want bm25", r) + } + time.Sleep(200 * time.Millisecond) + if got := f.embedded.Load(); got != 0 { + t.Errorf("embedded %d inputs with no cache to keep them in", got) + } +} + +// A document that fits the digest budget whole is sent whole, as today. +func TestSmallDocumentIsWhole(t *testing.T) { + f := newFakeDixie(t, false) + e := f.engine(t, NewGate(2, 8, 5*time.Second), true) + r := e.Run(context.Background(), "--- page 1 ---\n\nOne short page.", "what is it?", "https://ex.test/e.pdf", dcache.ClassFleet) + if r.Fallback != "" || r.Retrieval != "whole" { + t.Fatalf("%+v, want whole", r) + } +}