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

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)
	}
}