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 <ornith-1.5-35b-a3b@llama.cpp.invalid>
2026-09-26 18:44:28 -07:00
parent 7595ad8519
commit f6644c9278
6 changed files with 590 additions and 57 deletions
@@ -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)
+	}
+}