From 7595ad85192bf273bb9d4b3f96b67387ef52a9ec Mon Sep 17 00:00:00 2001 From: Kyle Isom Date: Sat, 26 Sep 2026 17:33:50 -0700 Subject: [PATCH] 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)) + } +}