From 7595ad85192bf273bb9d4b3f96b67387ef52a9ec Mon Sep 17 00:00:00 2001 From: Kyle Isom Date: Sat, 26 Sep 2026 17:33:50 -0700 Subject: [PATCH 1/4] digest: SplitPages keeps every chunk on its page; BM25 scorer SplitPages is Split's budget/line splitting 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, so every chunk names its page. BM25 (k1=1.2, b=0.75) ranks chunks lexically with counted term frequencies and ln(1+(N-df+0.5)/(df+0.5)) IDF. Co-Authored-By: ornith-1.5-35b-a3b --- docs/implementer-log.md | 1 + internal/digest/bm25.go | 106 +++++++++++++++++++++++++++++ internal/digest/bm25_test.go | 89 ++++++++++++++++++++++++ internal/digest/splitpages.go | 74 ++++++++++++++++++++ internal/digest/splitpages_test.go | 90 ++++++++++++++++++++++++ 5 files changed, 360 insertions(+) create mode 100644 internal/digest/bm25.go create mode 100644 internal/digest/bm25_test.go create mode 100644 internal/digest/splitpages.go create mode 100644 internal/digest/splitpages_test.go diff --git a/docs/implementer-log.md b/docs/implementer-log.md index 45ef4f6..ee49493 100644 --- a/docs/implementer-log.md +++ b/docs/implementer-log.md @@ -11,6 +11,7 @@ One row per task session that ends in a commit. Newest last. | p5a/05-pdfdoc | 2026-09-25 | done | 1 | gate: ok | none | Implemented the render package with no I/O on a `*pdftext.Doc`. `Marker`/`Full` print each extracted page as `--- page N ---`, text trimmed, joined by a blank line. `Render` dispatches grep → page range → whole. Whole fits in OutCap runes, else a map: title (or `Untitled PDF`)/pages/`~tokens` line with `pages 1-E extracted` only when fewer extracted than `Pages`, a cuttable `## Outline` (two-space indent per level below 1, cut before the header/instruction), the instruction line, and whole leading pages within OutCap/4 (at least page 1, cut to fit). Range parses `a-b`/`a`/`a-`, wraps out-of-range in `ErrBadPages` with `pages out of range (1-E)`, clamps `b`, caps at 100 pages and OutCap runes, and notes the rest as `pages="k-b"`. Grep is RE2 with smart-case `(?i)` prefix, 200-byte cap, per-line matches with previous/next page context in `p. P-`/`p. P:` form, first 100 shown, total counted; RE2 keeps the `(a+)+$` test linear. Every length is in runes. Given test byte-identical; make gate ok. | | 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. | | 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/bm25.go b/internal/digest/bm25.go new file mode 100644 index 0000000..610427e --- /dev/null +++ b/internal/digest/bm25.go @@ -0,0 +1,106 @@ +package digest + +import ( + "math" + "strings" + "unicode" +) + +const ( + bm25K1 = 1.2 + bm25B = 0.75 + bm25Min = 2 +) + +// BM25 is a lexical ranker over a fixed set of documents. Terms are the same tokens as +// terms() — lower-case runs of letters and digits, at least bm25Min characters — counted +// with frequency rather than kept as a set. +type BM25 struct { + nDocs int + avgDL float64 + df map[string]int + toks []map[string]int + len []int +} + +// NewBM25 indexes docs with the default BM25 parameters (k1 = 1.2, b = 0.75). +func NewBM25(docs []string) *BM25 { + b := &BM25{ + nDocs: len(docs), + df: map[string]int{}, + toks: make([]map[string]int, len(docs)), + len: make([]int, len(docs)), + } + var total int + for i, doc := range docs { + counts := docTerms(doc) + n := 0 + for _, c := range counts { + n += c + } + b.toks[i] = counts + b.len[i] = n + total += n + for t := range counts { + b.df[t]++ + } + } + if len(docs) > 0 { + b.avgDL = float64(total) / float64(len(docs)) + } + return b +} + +// 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 { + return make([]float64, len(b.toks)) + } + out := make([]float64, len(b.toks)) + qterms := docTerms(query) + if len(qterms) == 0 { + return out + } + avgDL := b.avgDL + if avgDL == 0 { + avgDL = 1 + } + for i, toks := range b.toks { + var score float64 + for term := range qterms { + tf, ok := toks[term] + if !ok { + continue + } + idf := math.Log(1 + (float64(b.nDocs)-float64(b.df[term])+0.5)/(float64(b.df[term])+0.5)) + denom := float64(tf) + bm25K1*(1-bm25B+bm25B*float64(b.len[i])/avgDL) + score += idf * (float64(tf) * (bm25K1 + 1)) / denom + } + out[i] = score + } + return out +} + +// docTerms counts occurrences of each BM25 term in s, using the same tokenization as terms(). +func docTerms(s string) map[string]int { + out := map[string]int{} + var b strings.Builder + flush := func() { + t := b.String() + b.Reset() + if len(t) < bm25Min { + return + } + out[t]++ + } + for _, r := range strings.ToLower(s) { + if unicode.IsLetter(r) || unicode.IsDigit(r) { + b.WriteRune(r) + } else { + flush() + } + } + flush() + return out +} diff --git a/internal/digest/bm25_test.go b/internal/digest/bm25_test.go new file mode 100644 index 0000000..ca9f084 --- /dev/null +++ b/internal/digest/bm25_test.go @@ -0,0 +1,89 @@ +package digest + +// P5b task 01: BM25, the lexical ranking used while a long document's embeddings are still +// being built (and inside the hybrid score after). + +import ( + "math" + "strings" + "testing" +) + +func TestBM25RanksTheRareTerm(t *testing.T) { + docs := []string{ + "the ship sailed from helsinki to tallinn in the morning", + "the ship stayed in port all day", + "a zanzibarite sample was found near the harbour", + "the harbour master logged the ship", + } + s := NewBM25(docs).Scores("where was the zanzibarite found") + if len(s) != len(docs) { + t.Fatalf("%d scores for %d docs", len(s), len(docs)) + } + best := 0 + for i := range s { + if s[i] > s[best] { + best = i + } + } + if best != 2 { + t.Errorf("best = doc %d (%v), want doc 2, the only one with the rare term", best, s) + } +} + +func TestBM25Properties(t *testing.T) { + filler := strings.Repeat("word ", 50) + docs := []string{ + "alpha " + filler, // 0: tf 1 + strings.Repeat("alpha ", 10) + filler, // 1: tf 10, same filler + "alpha common", // 2: tf 1, short + "common " + filler, // 3: no alpha + } + idx := NewBM25(docs) + s := idx.Scores("alpha") + if !(s[1] > s[0]) { + t.Errorf("more occurrences must score higher: tf10 %v vs tf1 %v", s[1], s[0]) + } + if !(s[1] < 10*s[0]) { + t.Errorf("term frequency must saturate: tf10 %v, 10 × tf1 %v", s[1], 10*s[0]) + } + if !(s[2] > s[0]) { + t.Errorf("the same term in a shorter document must score higher: short %v, long %v", s[2], s[0]) + } + if s[3] != 0 { + t.Errorf("a document without the term scores %v, want 0", s[3]) + } + // IDF: a term in every document is worth less than a rare one. + all := NewBM25([]string{"xx yy", "xx zz", "xx ww"}).Scores("xx") + rare := NewBM25([]string{"xx yy", "qq zz", "qq ww"}).Scores("xx") + if !(rare[0] > all[0]) { + t.Errorf("idf: a term in 1 of 3 docs must outscore one in 3 of 3: %v vs %v", rare[0], all[0]) + } +} + +func TestBM25Degenerate(t *testing.T) { + for _, docs := range [][]string{nil, {}, {""}, {"", "a b"}} { + s := NewBM25(docs).Scores("a") + if len(s) != len(docs) { + t.Errorf("%q: %d scores", docs, len(s)) + } + for _, v := range s { + if math.IsNaN(v) || math.IsInf(v, 0) || v < 0 { + t.Errorf("%q: score %v", docs, v) + } + } + } + for _, v := range NewBM25([]string{"a b", "c d"}).Scores("") { + if v != 0 { + t.Errorf("empty query scored %v", v) + } + } + // Deterministic. + a := NewBM25([]string{"alpha beta", "beta gamma", "gamma alpha"}).Scores("alpha gamma") + b := NewBM25([]string{"alpha beta", "beta gamma", "gamma alpha"}).Scores("alpha gamma") + for i := range a { + if a[i] != b[i] { + t.Fatal("scores differ between identical runs") + } + } +} diff --git a/internal/digest/splitpages.go b/internal/digest/splitpages.go new file mode 100644 index 0000000..f4406d9 --- /dev/null +++ b/internal/digest/splitpages.go @@ -0,0 +1,74 @@ +package digest + +import ( + "fmt" + "regexp" + "strconv" + "strings" +) + +// pageMarker matches a page header line like "--- page 7 ---"; the number is group 1. +var pageMarker = regexp.MustCompile(`^--- page (\d+) ---$`) + +// SplitPages splits text into chunks of at most maxTok tokens (EstTokens) on line +// boundaries, like Split, but page-aware. The text is expected to carry "--- page N ---" +// markers (as pdfdoc.Full produces). A marker line sets the current page; a chunk that +// starts after the first marker but not on a marker line is prefixed with +// "--- page N (cont.) ---\n\n" for the page its first line belongs to, so every chunk +// names the page it cites. Text with no marker lines yields exactly Split's result. +func SplitPages(text string, maxTok int) []string { + if maxTok <= 0 { + maxTok = ChunkTokens + } + lines := strings.Split(text, "\n") + var chunks []string + var buf strings.Builder + var firstLine string + var chunkPage int + hasLines, startedAfterMarker, markerSeen := false, false, false + currentPage := 0 + + flush := func() { + s := strings.TrimSpace(buf.String()) + if s == "" { + buf.Reset() + hasLines = false + return + } + if pageMarker.MatchString(firstLine) || !startedAfterMarker { + chunks = append(chunks, s) + } else { + chunks = append(chunks, fmt.Sprintf("--- page %d (cont.) ---\n\n", chunkPage)+s) + } + buf.Reset() + hasLines = false + } + + for _, line := range lines { + if m := pageMarker.FindStringSubmatch(line); m != nil { + if n, err := strconv.Atoi(m[1]); err == nil { + currentPage = n + markerSeen = true + } + } + next := line + "\n" + if hasLines { + est := EstTokens(buf.String() + next) + if startedAfterMarker { + est += EstTokens(fmt.Sprintf("--- page %d (cont.) ---\n\n", chunkPage)) + } + if est > maxTok { + flush() + } + } + if !hasLines { + firstLine = line + chunkPage = currentPage + startedAfterMarker = markerSeen + hasLines = true + } + buf.WriteString(next) + } + flush() + return chunks +} diff --git a/internal/digest/splitpages_test.go b/internal/digest/splitpages_test.go new file mode 100644 index 0000000..785698b --- /dev/null +++ b/internal/digest/splitpages_test.go @@ -0,0 +1,90 @@ +package digest + +// P5b task 01: SplitPages keeps every chunk of a PDF's text attached to its page, so a digest +// built from a chunk that starts mid-page can still cite the page. + +import ( + "fmt" + "regexp" + "strings" + "testing" +) + +// pagesText builds n pages of about perPage tokens each, every line naming its page. +func pagesText(n, perPage int) string { + var b strings.Builder + 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*12 < perPage; l++ { + fmt.Fprintf(&b, "page %d line %d alpha bravo charlie delta echo foxtrot golf hotel\n", p, l) + } + } + return b.String() +} + +var markerRE = regexp.MustCompile(`^--- page (\d+)( \(cont\.\))? ---`) + +func TestSplitPagesEveryChunkNamesItsPage(t *testing.T) { + text := pagesText(4, 3000) // each page ~3000 tokens: several chunks per page at 1200 + chunks := SplitPages(text, ChunkTokens) + if len(chunks) < 8 { + t.Fatalf("%d chunks, want at least 8 for 4 × 3000 tokens at %d", len(chunks), ChunkTokens) + } + cont := 0 + for i, c := range chunks { + m := markerRE.FindStringSubmatch(c) + if m == nil { + t.Fatalf("chunk %d does not start with a page marker: %.80q", i, c) + } + if m[2] != "" { + cont++ + // A continuation marker names the page the chunk's first line belongs to. + first := strings.SplitN(strings.TrimPrefix(c, m[0]), "\n", 3) + line := strings.TrimSpace(strings.Join(first, " ")) + if !strings.Contains(line, "page "+m[1]+" line ") { + t.Errorf("chunk %d marked page %s but starts with %.60q", i, m[1], line) + } + } + if EstTokens(c) > ChunkTokens+20 { + t.Errorf("chunk %d is %d tokens, over the %d budget (plus a marker line)", i, EstTokens(c), ChunkTokens) + } + } + if cont == 0 { + t.Error("no continuation markers: pages longer than a chunk must produce some") + } +} + +// Nothing is lost or duplicated: removing the continuation markers gives back the original lines. +func TestSplitPagesPreservesContent(t *testing.T) { + text := pagesText(3, 2500) + var got []string + for _, c := range SplitPages(text, ChunkTokens) { + for _, line := range strings.Split(c, "\n") { + if strings.HasSuffix(line, " (cont.) ---") || strings.TrimSpace(line) == "" { + continue + } + got = append(got, line) + } + } + var want []string + for _, line := range strings.Split(text, "\n") { + if strings.TrimSpace(line) != "" { + want = append(want, line) + } + } + if strings.Join(got, "\n") != strings.Join(want, "\n") { + t.Errorf("content changed: %d lines back, %d in", len(got), len(want)) + } +} + +// Text with no markers splits exactly as Split does. +func TestSplitPagesWithoutMarkersIsSplit(t *testing.T) { + text := strings.Repeat("plain words without any page markers at all\n", 2000) + a, b := SplitPages(text, ChunkTokens), Split(text, ChunkTokens) + if strings.Join(a, "\x00") != strings.Join(b, "\x00") { + t.Errorf("SplitPages differs from Split on unmarked text: %d vs %d chunks", len(a), len(b)) + } +} From f6644c92783e4b9303e7b84cb60ad379d5834cce Mon Sep 17 00:00:00 2001 From: Kyle Isom Date: Sat, 26 Sep 2026 18:44:28 -0700 Subject: [PATCH 2/4] digest: cache chunk embeddings; long documents answer from BM25 while embedding in the background Cached chunk embeddings change what a long document answers with: BM25 first, embedding uncached chunks (or the whole document) in the background so the next question on the same doc scores hybrid. selectText now returns its mode and the number of chunks that had a vector at answer-build time; the digest cache stores only hybrid answers, never a provisional bm25 one. The background worker batches through a Dixie gate and WaitBackground polls for true idle. Co-Authored-By: ornith-1.5-35b-a3b --- docs/implementer-log.md | 1 + internal/digest/background.go | 174 +++++++++++++++++++++ internal/digest/digest.go | 76 +++------ internal/digest/digest_test.go | 2 +- internal/digest/retrieval.go | 147 ++++++++++++++++++ internal/digest/retrieval_test.go | 247 ++++++++++++++++++++++++++++++ 6 files changed, 590 insertions(+), 57 deletions(-) create mode 100644 internal/digest/background.go create mode 100644 internal/digest/retrieval.go create mode 100644 internal/digest/retrieval_test.go 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) + } +} From 8cdb3e8537be606719160cf3a9c76c9e3b3b16e8 Mon Sep 17 00:00:00 2001 From: Kyle Isom Date: Sat, 26 Sep 2026 18:52:28 -0700 Subject: [PATCH 3/4] api: questions on long PDFs are answered (retrieval in the digest meta); README and skill The digest now handles long documents itself (task 02), so the API's refusal limit is only a safety cap: DefaultPDFMaxQuestionChunks 200 -> 5000 (the 422 message text is unchanged). In v1fetch the digest-meta object gains "retrieval" and "embedded" next to chunks_used/chunks_total, only when the digest ran without a fallback. The PDF section explains a long-doc question is answered first from BM25 while its chunks embed in the background and hybrid on later questions; skill/SKILL.md gains one sentence. Given tests byte-identical. Co-Authored-By: ornith-1.5-35b-a3b --- README.md | 6 ++- docs/implementer-log.md | 1 + internal/api/api.go | 2 + internal/api/pdf.go | 8 ++-- internal/api/pdf_retrieval_test.go | 75 ++++++++++++++++++++++++++++++ internal/api/pdf_test.go | 5 +- skill/SKILL.md | 3 +- 7 files changed, 92 insertions(+), 8 deletions(-) create mode 100644 internal/api/pdf_retrieval_test.go diff --git a/README.md b/README.md index e0775eb..680377e 100644 --- a/README.md +++ b/README.md @@ -87,8 +87,10 @@ Answers (through `/v1/fetch` or `/v2/scrape`): - `--question "…"` for a **digest** that, on a PDF, cites each fact as `(p. N)`. A range returns at most 100 pages; grep shows at most 100 matches and takes a pattern of at most -200 bytes. `pages` and `grep` on anything but a PDF are a `400`. A question over a document too long -to digest (more than 200 chunks) is refused with +200 bytes. `pages` and `grep` on anything but a PDF are a `400`. A question on a long document is +answered first from BM25 (`retrieval: "bm25"`) — lexical, before the document's chunks are +embedded in the background — and hybrid on later questions, once those embeddings exist. A question +over a document too long to answer even by retrieval (more than 5000 chunks) is refused with `422 document too long for question (N chunks); use grep or pages`. Scanned PDFs have no text layer and return `no text layer (scanned PDF?)`; point them at an OCR diff --git a/docs/implementer-log.md b/docs/implementer-log.md index db8dfed..f6fcb3b 100644 --- a/docs/implementer-log.md +++ b/docs/implementer-log.md @@ -14,6 +14,7 @@ One row per task session that ends in a commit. Newest last. | 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. | +| p5b/03-api | 2026-09-26 | done | 1 | gate: ok | none | The digest now answers long documents itself (task 02), so the API's refusal limit is only a safety cap: `DefaultPDFMaxQuestionChunks` 200→5000 in `internal/api/pdf.go` (the 422 message text unchanged, still "document too long for question (N chunks); use grep or pages"). In `v1fetch` (`internal/api/api.go`) the digest-meta map gains `"retrieval": dr.Retrieval` and `"embedded": dr.Embedded` next to `chunks_used`/`chunks_total`, only when the digest ran without a fallback. Copied the two given tests byte-identical from the plan — the replacement `pdf_test.go` is now the protected one and asserts `DefaultPDFMaxQuestionChunks == 5000`; `pdf_retrieval_test.go` asserts the first long-doc question returns `retrieval=bm25`, `embedded=0`. README PDF section: a long-doc question is answered first from BM25 (`retrieval: "bm25"`) while its chunks embed in the background and hybrid on later questions, plus the 5000-chunk cap. One sentence added to `skill/SKILL.md`. All api tests pass under `-race -count=3`; make gate ok; given files byte-identical. | ## P5a review (owner, 2026-09-26) diff --git a/internal/api/api.go b/internal/api/api.go index 715a057..a9cfe46 100644 --- a/internal/api/api.go +++ b/internal/api/api.go @@ -569,6 +569,8 @@ func (s *Server) v1fetch(w http.ResponseWriter, r *http.Request) { "model": dr.Model, "chunks_used": dr.ChunksUsed, "chunks_total": dr.ChunksTotal, + "retrieval": dr.Retrieval, + "embedded": dr.Embedded, "fallback": false, "cached": dr.Cached, } diff --git a/internal/api/pdf.go b/internal/api/pdf.go index 7c7179f..d13e313 100644 --- a/internal/api/pdf.go +++ b/internal/api/pdf.go @@ -20,9 +20,11 @@ import ( const ( // DefaultPDFTTL is how long a cached PDF is served before it is revalidated. DefaultPDFTTL = 24 * time.Hour - // DefaultPDFMaxQuestionChunks is the longest document (in digest chunks) a question is - // answered over; a longer one is refused in favour of grep or pages. - DefaultPDFMaxQuestionChunks = 200 + // DefaultPDFMaxQuestionChunks is the safety cap on how long a document a question is + // answered over; beyond it the question is refused in favour of grep or pages. Longer + // documents are handled by retrieval (BM25 first, embedding their chunks in the + // background), so this only bounds the un-answerable. + DefaultPDFMaxQuestionChunks = 5000 ) // runner returns the configured extraction child, or a zero Runner that runs diff --git a/internal/api/pdf_retrieval_test.go b/internal/api/pdf_retrieval_test.go new file mode 100644 index 0000000..2793160 --- /dev/null +++ b/internal/api/pdf_retrieval_test.go @@ -0,0 +1,75 @@ +package api + +// P5b task 03: a question on a long PDF is answered (by BM25 first) instead of refused, and the +// envelope says how the digest was built. + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "git.wntrmute.dev/kyle/sift/internal/digest" + "git.wntrmute.dev/kyle/sift/internal/pdftest" +) + +// fakeModel answers /embed with a fixed vector and /chat with a one-word SSE reply. +func fakeModel(t *testing.T) *httptest.Server { + mux := http.NewServeMux() + mux.HandleFunc("/embed", func(w http.ResponseWriter, r *http.Request) { + var body struct { + Input []string `json:"input"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + var data []map[string]any + for i := range body.Input { + data = append(data, map[string]any{"index": i, "embedding": []float64{1, float64(i % 7)}}) + } + _ = json.NewEncoder(w).Encode(map[string]any{"data": data}) + }) + mux.HandleFunc("/chat", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"answer (p. 3)"}}]}` + "\n\ndata: [DONE]\n\n")) + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv +} + +func TestPDFLongQuestionIsAnswered(t *testing.T) { + r := newPDFRig(t) + m := fakeModel(t) + db := r.s.Digest.Cache + r.s.Digest = &digest.Engine{ChatURL: m.URL + "/chat", EmbedURL: m.URL + "/embed", Cache: db, Gate: digest.NewGate(2, 8, 5*time.Second), Client: m.Client()} + var pages []string + for p := 1; p <= 260; p++ { // ~260 pages × 40 lines: well over EmbedSyncMax chunks + pages = append(pages, pdftest.Lorem(p, 40)) + } + o := newOrigin(t, "application/pdf", pdftest.Build(pdftest.Doc{Title: "Long", Pages: pages})) + code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "question": "what is on page three?", "context": false}) + if code != 200 { + t.Fatalf("long question: %d %v, want 200 (P5a refused this)", code, out) + } + dg, _ := out["digest"].(map[string]any) + if dg == nil || dg["retrieval"] != "bm25" || dg["embedded"] != 0.0 { + t.Errorf("digest meta = %v, want retrieval bm25 and embedded 0 on the first question", dg) + } + if c, _ := out["content"].(string); !strings.Contains(c, "answer") { + t.Errorf("content %q", c) + } +} + +func TestPDFSmallQuestionReportsRetrieval(t *testing.T) { + r := newPDFRig(t) + m := fakeModel(t) + db := r.s.Digest.Cache + r.s.Digest = &digest.Engine{ChatURL: m.URL + "/chat", EmbedURL: m.URL + "/embed", Cache: db, Gate: digest.NewGate(2, 8, 5*time.Second), Client: m.Client()} + o := newOrigin(t, "application/pdf", threePDF()) + code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "question": "what is on page two?", "context": false}) + dg, _ := out["digest"].(map[string]any) + if code != 200 || dg == nil || dg["retrieval"] != "whole" { + t.Errorf("small document: %d digest %v, want retrieval whole", code, dg) + } +} diff --git a/internal/api/pdf_test.go b/internal/api/pdf_test.go index 041978a..2874eea 100644 --- a/internal/api/pdf_test.go +++ b/internal/api/pdf_test.go @@ -396,7 +396,8 @@ func TestPDFQuestionTooLong(t *testing.T) { if msg, _ := out["error"].(string); code != 422 || !strings.Contains(msg, "too long") || !strings.Contains(msg, "grep") { t.Errorf("long question: %d %v; want 422 … too long … grep", code, out) } - if DefaultPDFMaxQuestionChunks != 200 { - t.Errorf("DefaultPDFMaxQuestionChunks = %d, want 200", DefaultPDFMaxQuestionChunks) + // P5b: long documents are answered by retrieval now; the limit is only a safety cap. + if DefaultPDFMaxQuestionChunks != 5000 { + t.Errorf("DefaultPDFMaxQuestionChunks = %d, want 5000", DefaultPDFMaxQuestionChunks) } } diff --git a/skill/SKILL.md b/skill/SKILL.md index 4a0fe97..7ea19c4 100644 --- a/skill/SKILL.md +++ b/skill/SKILL.md @@ -61,7 +61,8 @@ PDFs work: `sift fetch` extracts them and returns the text. A short one comes ba long one comes back as a map — an outline with page numbers, then the first pages. Then: - `sift fetch URL --pages 41-60` for a page range, `--grep 'pattern'` for matching lines, or - `--question "…"` for a digest that cites each fact as `(p. N)`. + `--question "…"` for a digest that cites each fact as `(p. N)`; `--question` works on long + PDFs, and the first answer may be `bm25` (lexical) — ask again later for a sharper one. - Scanned PDFs have no text layer: sift reports `no text layer (scanned PDF?)` — point it at an OCR source rather than trying to fix it here. - A range returns at most 100 pages and grep at most 100 matches. On a very long document From 64d08a358d44938e194fde4ca364b90f13f503b6 Mon Sep 17 00:00:00 2001 From: Kyle Isom Date: Sat, 26 Sep 2026 19:02:51 -0700 Subject: [PATCH 4/4] 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) + } +}