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>
This commit is contained in:
2026-09-26 18:44:28 -07:00
co-authored by ornith-1.5-35b-a3b
parent 7595ad8519
commit f6644c9278
6 changed files with 590 additions and 57 deletions
+1
View File
@@ -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/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. | | 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/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/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) ## P5a review (owner, 2026-09-26)
+174
View File
@@ -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)
}
}
+20 -56
View File
@@ -33,6 +33,8 @@ const (
ChunkTokens = 1200 ChunkTokens = 1200
EmbedBatch = 16 EmbedBatch = 16
EmbedMaxRunes = 3000 EmbedMaxRunes = 3000
EmbedSyncMax = 48
MaxDocChunks = 5000
DefaultPermit = 2 DefaultPermit = 2
DefaultQDepth = 8 DefaultQDepth = 8
DefaultQWait = 20 * time.Second DefaultQWait = 20 * time.Second
@@ -191,6 +193,9 @@ type Engine struct {
Log *slog.Logger Log *slog.Logger
Cache *dcache.DB Cache *dcache.DB
bgOnce sync.Once
bg *bgWorker
mu sync.Mutex mu sync.Mutex
Calls uint64 Calls uint64
Fallback map[string]uint64 Fallback map[string]uint64
@@ -205,6 +210,8 @@ type Result struct {
DigestMS int64 DigestMS int64
ReturnedTok int ReturnedTok int
Cached bool 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 { 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.ChunksUsed = h.ChunksUsed
r.ReturnedTok = EstTokens(h.Text) r.ReturnedTok = EstTokens(h.Text)
r.Cached = true r.Cached = true
r.Retrieval = "cached"
r.DigestMS = time.Since(start).Milliseconds() r.DigestMS = time.Since(start).Milliseconds()
return r return r
} }
@@ -327,8 +335,10 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL, callerCla
dctx, cancel := context.WithTimeout(ctx, e.timeout()) dctx, cancel := context.WithTimeout(ctx, e.timeout())
defer cancel() 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.ChunksUsed, r.ChunksTotal = nUsed, nTotal
r.Retrieval = mode
r.Embedded = embedded
if err != nil { if err != nil {
e.logFallback("embed", err) e.logFallback("embed", err)
r.Fallback = "embed" r.Fallback = "embed"
@@ -359,72 +369,26 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL, callerCla
} }
r.Text = text r.Text = text
r.ReturnedTok = EstTokens(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) e.Cache.PutDigest(dkey, text, r.ChunksUsed)
} }
return r 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 { 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 { if len(chunks) == 0 {
return cleaned, 1, 1, nil return cleaned, 1, 1, "whole", 0, nil
} }
if question == "" { 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)...)) return e.hybridOrBM25(ctx, chunks, question, contentHash, callerClass)
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
} }
func capEmbed(chunks []string) []string { func capEmbed(chunks []string) []string {
+1 -1
View File
@@ -374,7 +374,7 @@ func TestPinChunk0Infobox(t *testing.T) {
if EstTokens(cleaned) <= SmallDoc { if EstTokens(cleaned) <= SmallDoc {
t.Fatalf("fixture too small: %d", EstTokens(cleaned)) 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 { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
+147
View File
@@ -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)
}
+247
View File
@@ -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)
}
}