P5b and P5c plans: long-document retrieval (BM25 first, background embedding) and /v1/extract + sift extract; given tests

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

docs/plans/p5b/_files/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)
+	}
+}

docs/plans/p5b/_files/internal/api/pdf_test.go

@@ -0,0 +1,403 @@
+package api
+
+// P5a task 06: PDFs through /v1/fetch and /v2/scrape — extraction in the child, the pool, the
+// document cache, the map, pages, grep, and the errors.
+//
+// The child is the test binary itself (TestMain below), as in internal/pdftext's tests.
+// SIFT_PDFTEXT_RUNLOG makes every child start append a line to a file, so a test can count
+// extractions; SIFT_PDFTEXT_SLEEP slows a child down. Both hooks live only in this file.
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"net/netip"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"git.wntrmute.dev/kyle/sift/internal/auth"
+	"git.wntrmute.dev/kyle/sift/internal/dcache"
+	"git.wntrmute.dev/kyle/sift/internal/digest"
+	"git.wntrmute.dev/kyle/sift/internal/fetch"
+	"git.wntrmute.dev/kyle/sift/internal/guard"
+	"git.wntrmute.dev/kyle/sift/internal/pdfpool"
+	"git.wntrmute.dev/kyle/sift/internal/pdftest"
+	"git.wntrmute.dev/kyle/sift/internal/pdftext"
+)
+
+func TestMain(m *testing.M) {
+	if os.Getenv("SIFT_PDFTEXT_CHILD") == "1" {
+		if p := os.Getenv("SIFT_PDFTEXT_RUNLOG"); p != "" {
+			f, err := os.OpenFile(p, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600)
+			if err == nil {
+				_, _ = f.WriteString("run\n")
+				_ = f.Close()
+			}
+		}
+		if d, err := time.ParseDuration(os.Getenv("SIFT_PDFTEXT_SLEEP")); err == nil {
+			time.Sleep(d)
+		}
+		os.Exit(pdftext.ChildMain(os.Args[1:], os.Stdin, os.Stdout))
+	}
+	os.Exit(m.Run())
+}
+
+func threePDF() []byte {
+	return pdftest.Build(pdftest.Doc{
+		Title:   "Three Pages",
+		Pages:   []string{"Chapter one begins here.", "The quick brown fox is on page two.", "Last page."},
+		Outline: []pdftest.Outline{{Title: "1 Intro", Page: 1, Level: 1}, {Title: "2 Fox", Page: 2, Level: 1}},
+	})
+}
+
+func longPDF(n int) []byte {
+	var pages []string
+	for p := 1; p <= n; p++ {
+		pages = append(pages, pdftest.Lorem(p, 40))
+	}
+	return pdftest.Build(pdftest.Doc{Title: "Long Thesis", Pages: pages, Outline: []pdftest.Outline{{Title: "1 Introduction", Page: 1, Level: 1}, {Title: "2 Method", Page: n / 2, Level: 1}}})
+}
+
+// origin serves one PDF and counts GETs. It honours If-None-Match against its current ETag.
+type origin struct {
+	mu   sync.Mutex
+	ct   string
+	body []byte
+	etag string
+	gets atomic.Int32
+	srv  *httptest.Server
+}
+
+func newOrigin(t *testing.T, ct string, body []byte) *origin {
+	o := &origin{ct: ct, body: body, etag: `"v1"`}
+	o.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		o.gets.Add(1)
+		o.mu.Lock()
+		defer o.mu.Unlock()
+		if r.Header.Get("If-None-Match") == o.etag {
+			w.WriteHeader(http.StatusNotModified)
+			return
+		}
+		w.Header().Set("ETag", o.etag)
+		if o.ct != "" {
+			w.Header().Set("Content-Type", o.ct)
+		}
+		_, _ = w.Write(o.body)
+	}))
+	t.Cleanup(o.srv.Close)
+	return o
+}
+
+func (o *origin) replace(body []byte, etag string) {
+	o.mu.Lock()
+	o.body, o.etag = body, etag
+	o.mu.Unlock()
+}
+
+type pdfRig struct {
+	s      *Server
+	fleet  string // bearer key of caller "test" (hermes-fleet)
+	other  string // bearer key of caller "other" (hermes-fleet)
+	gs     string // bearer key of restricted caller "straylight-gs"
+	runlog string
+	logs   *recHandler
+}
+
+func newPDFRig(t *testing.T, childEnv ...string) *pdfRig {
+	t.Helper()
+	dir := t.TempDir()
+	keys := filepath.Join(dir, "keys.toml")
+	key := func(name, ws string) string {
+		out := filepath.Join(dir, name)
+		if err := auth.Add(keys, name, ws, out); err != nil {
+			t.Fatal(err)
+		}
+		b, err := os.ReadFile(out)
+		if err != nil {
+			t.Fatal(err)
+		}
+		return strings.TrimSpace(string(b))
+	}
+	r := &pdfRig{fleet: key("test", "hermes-fleet"), other: key("other", "hermes-fleet"), gs: key("straylight-gs", "guardian-star"), runlog: filepath.Join(dir, "runs"), logs: &recHandler{}}
+	store, err := auth.Load(keys)
+	if err != nil {
+		t.Fatal(err)
+	}
+	g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}}
+	s := New("test", g, &fetch.Fetcher{Guard: g, UA: "sift-test"}, store, slog.New(r.logs))
+	s.PermitRestricted = []string{"straylight-gs"}
+	db, err := dcache.Open(filepath.Join(dir, "cache.db"), 0, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = db.Close() })
+	s.Digest = &digest.Engine{Cache: db} // the cache lives on the digest engine; no model is configured
+	s.PDFRunner = &pdftext.Runner{Exe: os.Args[0], Env: append([]string{"SIFT_PDFTEXT_CHILD=1", "SIFT_PDFTEXT_RUNLOG=" + r.runlog}, childEnv...), Deadline: 20 * time.Second}
+	r.s = s
+	return r
+}
+
+func (r *pdfRig) runs() int {
+	b, _ := os.ReadFile(r.runlog)
+	return strings.Count(string(b), "run\n")
+}
+
+func (r *pdfRig) fetch(t *testing.T, key string, body map[string]any) (int, map[string]any, http.Header) {
+	t.Helper()
+	rec := postJSON(t, r.s, "/v1/fetch", key, body)
+	var out map[string]any
+	_ = json.Unmarshal(rec.Body.Bytes(), &out)
+	return rec.Code, out, rec.Header()
+}
+
+func pdfInfo(t *testing.T, out map[string]any) map[string]any {
+	t.Helper()
+	p, ok := out["pdf"].(map[string]any)
+	if !ok {
+		t.Fatalf("envelope has no pdf object: %v", out)
+	}
+	return p
+}
+
+func TestPDFCleanFetch(t *testing.T) {
+	for _, ct := range []string{"application/pdf", "application/octet-stream", ""} {
+		r := newPDFRig(t)
+		o := newOrigin(t, ct, threePDF())
+		code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
+		if code != 200 {
+			t.Fatalf("%q: %d %v", ct, code, out)
+		}
+		content, _ := out["content"].(string)
+		if !strings.Contains(content, "--- page 2 ---") || !strings.Contains(content, "quick brown fox") || out["mode"] != "clean" {
+			t.Errorf("%q: content %q mode %v", ct, content, out["mode"])
+		}
+		p := pdfInfo(t, out)
+		if p["pages"] != 3.0 || p["title"] != "Three Pages" || p["map"] != false || p["first"] != 1.0 || p["last"] != 3.0 {
+			t.Errorf("%q: pdf = %v", ct, p)
+		}
+		if ol, _ := p["outline"].([]any); len(ol) != 2 {
+			t.Errorf("%q: outline = %v", ct, p["outline"])
+		}
+	}
+}
+
+func TestPDFMap(t *testing.T) {
+	r := newPDFRig(t)
+	r.s.OutCap = 20000
+	o := newOrigin(t, "application/pdf", longPDF(40)) // ~40 × 2,700 characters
+	code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
+	if code != 200 {
+		t.Fatalf("%d %v", code, out)
+	}
+	content, _ := out["content"].(string)
+	p := pdfInfo(t, out)
+	if p["map"] != true || !strings.Contains(content, "## Outline") || !strings.Contains(content, "- 2 Method (p. 20)") || len([]rune(content)) > 20000 {
+		t.Errorf("map: pdf %v, %d chars:\n%.400s", p, len([]rune(content)), content)
+	}
+}
+
+func TestPDFPagesAndGrep(t *testing.T) {
+	r := newPDFRig(t)
+	o := newOrigin(t, "application/pdf", threePDF())
+	code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "pages": "2-3"})
+	content, _ := out["content"].(string)
+	if code != 200 || strings.Contains(content, "--- page 1 ---") || !strings.Contains(content, "--- page 3 ---") {
+		t.Errorf("pages 2-3: %d %q", code, content)
+	}
+	if p := pdfInfo(t, out); p["first"] != 2.0 || p["last"] != 3.0 {
+		t.Errorf("pages 2-3: pdf %v", p)
+	}
+	for _, bad := range []map[string]any{{"pages": "0-1"}, {"pages": "9"}, {"grep": strings.Repeat("a", 201)}, {"grep": "(unclosed"}} {
+		bad["url"] = o.srv.URL
+		if code, out, _ := r.fetch(t, r.fleet, bad); code != 400 || out["success"] != false {
+			t.Errorf("%v: %d %v, want 400", bad, code, out)
+		}
+	}
+	code, out, _ = r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "grep": "quick brown"})
+	content, _ = out["content"].(string)
+	if code != 200 || !strings.Contains(content, "p. 2: The quick brown fox is on page two.") || pdfInfo(t, out)["grep_hits"] != 1.0 {
+		t.Errorf("grep: %d %q %v", code, content, out["pdf"])
+	}
+	// pages and grep on a non-PDF are refused rather than ignored.
+	html := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "text/html")
+		_, _ = io.WriteString(w, "<html><body><p>hi</p></body></html>")
+	}))
+	defer html.Close()
+	if code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": html.URL, "pages": "1"}); code != 400 {
+		t.Errorf("pages on html: %d %v, want 400", code, out)
+	}
+}
+
+func TestPDFCache(t *testing.T) {
+	r := newPDFRig(t)
+	o := newOrigin(t, "application/pdf", threePDF())
+	get := func(extra map[string]any) string {
+		body := map[string]any{"url": o.srv.URL}
+		for k, v := range extra {
+			body[k] = v
+		}
+		code, out, _ := r.fetch(t, r.fleet, body)
+		if code != 200 {
+			t.Fatalf("%v: %d %v", extra, code, out)
+		}
+		c, _ := out["content"].(string)
+		return c
+	}
+	get(nil)
+	if o.gets.Load() != 1 || r.runs() != 1 {
+		t.Fatalf("first fetch: %d GETs, %d extractions; want 1 and 1", o.gets.Load(), r.runs())
+	}
+	if c := get(map[string]any{"pages": "2"}); !strings.Contains(c, "quick brown fox") {
+		t.Errorf("cached pages=2: %q", c)
+	}
+	if o.gets.Load() != 1 || r.runs() != 1 {
+		t.Errorf("fresh cache hit: %d GETs, %d extractions; want still 1 and 1", o.gets.Load(), r.runs())
+	}
+	// Stale: a conditional GET, answered 304, re-uses the text.
+	r.s.PDFTTL = time.Nanosecond
+	time.Sleep(2 * time.Millisecond)
+	get(map[string]any{"grep": "Chapter"})
+	if o.gets.Load() != 2 || r.runs() != 1 {
+		t.Errorf("after 304: %d GETs, %d extractions; want 2 and 1", o.gets.Load(), r.runs())
+	}
+	// The document changed: a full GET and a new extraction.
+	o.replace(pdftest.Build(pdftest.Doc{Pages: []string{"Brand new text."}}), `"v2"`)
+	if c := get(nil); !strings.Contains(c, "Brand new text.") {
+		t.Errorf("changed document: %q", c)
+	}
+	if o.gets.Load() != 3 || r.runs() != 2 {
+		t.Errorf("after change: %d GETs, %d extractions; want 3 and 2", o.gets.Load(), r.runs())
+	}
+	if DefaultPDFTTL != 24*time.Hour {
+		t.Errorf("DefaultPDFTTL = %v, want 24h", DefaultPDFTTL)
+	}
+}
+
+func TestPDFErrors(t *testing.T) {
+	r := newPDFRig(t)
+	enc, err := os.ReadFile(filepath.Join("..", "..", "tests", "testdata", "pdf", "encrypted-user.pdf"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, tc := range []struct {
+		name string
+		body []byte
+		want string
+		hint string
+	}{
+		{"scanned", pdftest.Build(pdftest.Doc{Pages: []string{"", ""}, Scanned: []int{1, 2}}), "no text layer", "ocr"},
+		{"encrypted", enc, "encrypted", ""},
+		{"garbage", []byte("%PDF-1.4\nthis is not really a pdf"), "unreadable", ""},
+	} {
+		o := newOrigin(t, "application/pdf", tc.body)
+		code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
+		msg, _ := out["error"].(string)
+		if code != 422 || !strings.Contains(msg, tc.want) || (tc.hint != "" && out["hint"] != tc.hint) {
+			t.Errorf("%s: %d %v; want 422 %q hint %q", tc.name, code, out, tc.want, tc.hint)
+		}
+	}
+}
+
+func TestPDFTimeout(t *testing.T) {
+	r := newPDFRig(t, "SIFT_PDFTEXT_SLEEP=10s")
+	r.s.PDFRunner.Deadline = 300 * time.Millisecond
+	o := newOrigin(t, "application/pdf", threePDF())
+	code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
+	if msg, _ := out["error"].(string); code != 422 || !strings.Contains(msg, "timed out") {
+		t.Errorf("slow child: %d %v; want 422 … timed out", code, out)
+	}
+}
+
+func TestPDFBusy(t *testing.T) {
+	r := newPDFRig(t, "SIFT_PDFTEXT_SLEEP=1500ms")
+	r.s.PDFPool = pdfpool.New(1, 0, 100*time.Millisecond)
+	a := newOrigin(t, "application/pdf", threePDF())
+	b := newOrigin(t, "application/pdf", longPDF(2))
+	done := make(chan int, 1)
+	go func() {
+		code, _, _ := r.fetch(t, r.fleet, map[string]any{"url": a.srv.URL})
+		done <- code
+	}()
+	time.Sleep(400 * time.Millisecond) // the first extraction holds the only slot
+	code, out, hdr := r.fetch(t, r.other, map[string]any{"url": b.srv.URL})
+	if code != 503 || hdr.Get("Retry-After") != "10" || !strings.Contains(fmt.Sprint(out["error"]), "busy") {
+		t.Errorf("second caller with the pool full: %d %v Retry-After %q; want 503, busy, 10", code, out, hdr.Get("Retry-After"))
+	}
+	code, out, _ = r.fetch(t, r.fleet, map[string]any{"url": b.srv.URL})
+	if code != 429 {
+		t.Errorf("same caller, second document while its first is running: %d %v; want 429", code, out)
+	}
+	if c := <-done; c != 200 {
+		t.Errorf("first fetch: %d", c)
+	}
+}
+
+func TestPDFRestrictedCaller(t *testing.T) {
+	r := newPDFRig(t)
+	o := newOrigin(t, "application/pdf", threePDF())
+	if code, out, _ := r.fetch(t, r.gs, map[string]any{"url": o.srv.URL}); code != 200 {
+		t.Fatalf("restricted fetch: %d %v", code, out)
+	}
+	r.logs.mu.Lock()
+	for _, rec := range r.logs.recs {
+		var b strings.Builder
+		b.WriteString(rec.Message)
+		rec.Attrs(func(a slog.Attr) bool { fmt.Fprintf(&b, " %s=%v", a.Key, a.Value); return true })
+		line := b.String()
+		if strings.Contains(line, "127.0.0.1") || strings.Contains(line, "Three Pages") || strings.Contains(line, "quick brown") {
+			t.Errorf("a restricted caller's PDF fetch logged %q", line)
+		}
+	}
+	r.logs.mu.Unlock()
+	// The fleet does not see the restricted caller's cached text: it extracts its own copy.
+	if code, _, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL}); code != 200 || r.runs() != 2 {
+		t.Errorf("fleet after restricted: %d, %d extractions; want 200 and 2 (separate cache partitions)", code, r.runs())
+	}
+}
+
+func TestPDFScrape(t *testing.T) {
+	r := newPDFRig(t)
+	o := newOrigin(t, "application/pdf", threePDF())
+	rec := postJSON(t, r.s, "/v2/scrape", r.fleet, map[string]any{"url": o.srv.URL, "formats": []string{"markdown"}})
+	var out struct {
+		Success bool `json:"success"`
+		Data    struct {
+			Markdown string         `json:"markdown"`
+			Metadata map[string]any `json:"metadata"`
+		} `json:"data"`
+	}
+	if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil || rec.Code != 200 || !out.Success {
+		t.Fatalf("%d %s", rec.Code, rec.Body.String())
+	}
+	if !strings.Contains(out.Data.Markdown, "--- page 2 ---") || out.Data.Metadata["numPages"] != 3.0 || out.Data.Metadata["title"] != "Three Pages" {
+		t.Errorf("scrape: markdown %q metadata %v", out.Data.Markdown, out.Data.Metadata)
+	}
+	if ct, _ := out.Data.Metadata["contentType"].(string); !strings.HasPrefix(ct, "application/pdf") {
+		t.Errorf("contentType = %q", ct)
+	}
+}
+
+// A question on a document longer than PDFMaxQuestionChunks is refused with a pointer to grep
+// and pages, before any model call.
+func TestPDFQuestionTooLong(t *testing.T) {
+	r := newPDFRig(t)
+	r.s.PDFMaxQuestionChunks = 2
+	o := newOrigin(t, "application/pdf", longPDF(12))
+	code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "question": "what is the method?"})
+	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)
+	}
+	// 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)
+	}
+}

docs/plans/p5b/_files/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")
+		}
+	}
+}

docs/plans/p5b/_files/internal/digest/retrieval_test.go

@@ -0,0 +1,244 @@
+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(12, 0) // ~12 pages, a few chunks: under EmbedSyncMax
+	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)
+	}
+}

docs/plans/p5b/_files/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))
+	}
+}