Merge P5b: long-document questions (page-aware chunks, BM25 first, cached embeddings built in the background)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
@@ -87,8 +87,10 @@ Answers (through `/v1/fetch` or `/v2/scrape`):
|
||||
- `--question "…"` for a **digest** that, on a PDF, cites each fact as `(p. N)`.
|
||||
|
||||
A range returns at most 100 pages; grep shows at most 100 matches and takes a pattern of at most
|
||||
200 bytes. `pages` and `grep` on anything but a PDF are a `400`. A question over a document too long
|
||||
to digest (more than 200 chunks) is refused with
|
||||
200 bytes. `pages` and `grep` on anything but a PDF are a `400`. A question on a long document is
|
||||
answered first from BM25 (`retrieval: "bm25"`) — lexical, before the document's chunks are
|
||||
embedded in the background — and hybrid on later questions, once those embeddings exist. A question
|
||||
over a document too long to answer even by retrieval (more than 5000 chunks) is refused with
|
||||
`422 document too long for question (N chunks); use grep or pages`.
|
||||
|
||||
Scanned PDFs have no text layer and return `no text layer (scanned PDF?)`; point them at an OCR
|
||||
|
||||
@@ -11,7 +11,10 @@ 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. |
|
||||
| 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. |
|
||||
| p5b/03-api | 2026-09-26 | done | 1 | gate: ok | none | The digest now answers long documents itself (task 02), so the API's refusal limit is only a safety cap: `DefaultPDFMaxQuestionChunks` 200→5000 in `internal/api/pdf.go` (the 422 message text unchanged, still "document too long for question (N chunks); use grep or pages"). In `v1fetch` (`internal/api/api.go`) the digest-meta map gains `"retrieval": dr.Retrieval` and `"embedded": dr.Embedded` next to `chunks_used`/`chunks_total`, only when the digest ran without a fallback. Copied the two given tests byte-identical from the plan — the replacement `pdf_test.go` is now the protected one and asserts `DefaultPDFMaxQuestionChunks == 5000`; `pdf_retrieval_test.go` asserts the first long-doc question returns `retrieval=bm25`, `embedded=0`. README PDF section: a long-doc question is answered first from BM25 (`retrieval: "bm25"`) while its chunks embed in the background and hybrid on later questions, plus the 5000-chunk cap. One sentence added to `skill/SKILL.md`. All api tests pass under `-race -count=3`; make gate ok; given files byte-identical. |
|
||||
|
||||
## P5a review (owner, 2026-09-26)
|
||||
|
||||
|
||||
@@ -569,6 +569,8 @@ func (s *Server) v1fetch(w http.ResponseWriter, r *http.Request) {
|
||||
"model": dr.Model,
|
||||
"chunks_used": dr.ChunksUsed,
|
||||
"chunks_total": dr.ChunksTotal,
|
||||
"retrieval": dr.Retrieval,
|
||||
"embedded": dr.Embedded,
|
||||
"fallback": false,
|
||||
"cached": dr.Cached,
|
||||
}
|
||||
|
||||
+5
-3
@@ -20,9 +20,11 @@ import (
|
||||
const (
|
||||
// DefaultPDFTTL is how long a cached PDF is served before it is revalidated.
|
||||
DefaultPDFTTL = 24 * time.Hour
|
||||
// DefaultPDFMaxQuestionChunks is the longest document (in digest chunks) a question is
|
||||
// answered over; a longer one is refused in favour of grep or pages.
|
||||
DefaultPDFMaxQuestionChunks = 200
|
||||
// DefaultPDFMaxQuestionChunks is the safety cap on how long a document a question is
|
||||
// answered over; beyond it the question is refused in favour of grep or pages. Longer
|
||||
// documents are handled by retrieval (BM25 first, embedding their chunks in the
|
||||
// background), so this only bounds the un-answerable.
|
||||
DefaultPDFMaxQuestionChunks = 5000
|
||||
)
|
||||
|
||||
// runner returns the configured extraction child, or a zero Runner that runs
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -396,7 +396,8 @@ func TestPDFQuestionTooLong(t *testing.T) {
|
||||
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)
|
||||
}
|
||||
if DefaultPDFMaxQuestionChunks != 200 {
|
||||
t.Errorf("DefaultPDFMaxQuestionChunks = %d, want 200", DefaultPDFMaxQuestionChunks)
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
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)
|
||||
|
||||
// Forget the document whether it finished or not: dedupe covers only queued and
|
||||
// in-progress work, so a later question re-queues a document whose embedding failed.
|
||||
w.mu.Lock()
|
||||
delete(w.seen, job.contentHash+"|"+job.callerClass)
|
||||
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]]
|
||||
}
|
||||
w.acquireGate()
|
||||
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 {
|
||||
if k < len(vecs) && len(vecs[k]) > 0 {
|
||||
e.Cache.PutEmbed(key(idx[b+k]), vecs[k])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// acquireGate takes the Dixie gate, sleeping a second after each refusal (queue full or wait
|
||||
// expired) so interactive digests get their turns. Background work has no deadline: it waits.
|
||||
func (w *bgWorker) acquireGate() {
|
||||
for {
|
||||
if ok, _ := w.engine.Gate.Acquire(context.Background()); ok {
|
||||
return
|
||||
}
|
||||
time.Sleep(time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
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 {
|
||||
return nil
|
||||
}
|
||||
if 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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
}
|
||||
+20
-56
@@ -33,6 +33,8 @@ const (
|
||||
ChunkTokens = 1200
|
||||
EmbedBatch = 16
|
||||
EmbedMaxRunes = 3000
|
||||
EmbedSyncMax = 48
|
||||
MaxDocChunks = 5000
|
||||
DefaultPermit = 2
|
||||
DefaultQDepth = 8
|
||||
DefaultQWait = 20 * time.Second
|
||||
@@ -191,6 +193,9 @@ type Engine struct {
|
||||
Log *slog.Logger
|
||||
Cache *dcache.DB
|
||||
|
||||
bgOnce sync.Once
|
||||
bg *bgWorker
|
||||
|
||||
mu sync.Mutex
|
||||
Calls uint64
|
||||
Fallback map[string]uint64
|
||||
@@ -205,6 +210,8 @@ type Result struct {
|
||||
DigestMS int64
|
||||
ReturnedTok int
|
||||
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 {
|
||||
@@ -305,6 +312,7 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL, callerCla
|
||||
r.ChunksUsed = h.ChunksUsed
|
||||
r.ReturnedTok = EstTokens(h.Text)
|
||||
r.Cached = true
|
||||
r.Retrieval = "cached"
|
||||
r.DigestMS = time.Since(start).Milliseconds()
|
||||
return r
|
||||
}
|
||||
@@ -327,8 +335,10 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL, callerCla
|
||||
dctx, cancel := context.WithTimeout(ctx, e.timeout())
|
||||
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.Retrieval = mode
|
||||
r.Embedded = embedded
|
||||
if err != nil {
|
||||
e.logFallback("embed", err)
|
||||
r.Fallback = "embed"
|
||||
@@ -359,72 +369,26 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL, callerCla
|
||||
}
|
||||
r.Text = 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)
|
||||
}
|
||||
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 {
|
||||
return cleaned, 1, 1, nil
|
||||
return cleaned, 1, 1, "whole", 0, nil
|
||||
}
|
||||
chunks := Split(cleaned, ChunkTokens)
|
||||
chunks := SplitPages(cleaned, ChunkTokens)
|
||||
if len(chunks) == 0 {
|
||||
return cleaned, 1, 1, nil
|
||||
return cleaned, 1, 1, "whole", 0, nil
|
||||
}
|
||||
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)...))
|
||||
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
|
||||
return e.hybridOrBM25(ctx, chunks, question, contentHash, callerClass)
|
||||
}
|
||||
|
||||
func capEmbed(chunks []string) []string {
|
||||
|
||||
@@ -374,7 +374,7 @@ func TestPinChunk0Infobox(t *testing.T) {
|
||||
if EstTokens(cleaned) <= SmallDoc {
|
||||
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 {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
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 {
|
||||
if len(vecs[j+1]) > 0 { // a vector the server left out is not cached as empty
|
||||
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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package digest
|
||||
|
||||
// P5b owner review: regression tests for the background embedder.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// flakyDixie is fakeDixie whose /embed answers 500 to its first failN requests.
|
||||
func flakyDixie(t *testing.T, failN int64) (*fakeDixie, *atomic.Int64) {
|
||||
f := &fakeDixie{}
|
||||
var calls atomic.Int64
|
||||
var disc atomic.Bool
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/embed", func(w http.ResponseWriter, r *http.Request) {
|
||||
if calls.Add(1) <= failN {
|
||||
http.Error(w, "boom", http.StatusInternalServerError)
|
||||
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(f.srv.Close)
|
||||
return f, &calls
|
||||
}
|
||||
|
||||
// A document whose background embedding failed is queued again by a later question; it must
|
||||
// not stay "seen" (and so BM25-only) for the life of the process.
|
||||
func TestBackgroundRequeuesAfterFailure(t *testing.T) {
|
||||
f, _ := flakyDixie(t, 1)
|
||||
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
|
||||
doc := longDoc(300, 250)
|
||||
n := len(SplitPages(doc, ChunkTokens))
|
||||
wait := func() {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
if err := e.WaitBackground(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
q := "who signed the zanzibarite manifest?"
|
||||
if r := e.Run(context.Background(), doc, q, "https://ex.test/f.pdf", "fleet"); r.Retrieval != "bm25" {
|
||||
t.Fatalf("first: %+v", r)
|
||||
}
|
||||
wait() // the first batch failed: the document was dropped
|
||||
if r := e.Run(context.Background(), doc, q, "https://ex.test/f.pdf", "fleet"); r.Retrieval != "bm25" {
|
||||
t.Fatalf("second: %+v, want bm25 (nothing embedded yet)", r)
|
||||
}
|
||||
wait() // re-queued by the second question, and embedded this time
|
||||
r := e.Run(context.Background(), doc, q, "https://ex.test/f.pdf", "fleet")
|
||||
if r.Retrieval != "hybrid" || r.Embedded != n {
|
||||
t.Fatalf("third: %+v, want hybrid with all %d embedded", r, n)
|
||||
}
|
||||
}
|
||||
|
||||
// The background worker waits for a busy gate; it does not give up on the document.
|
||||
func TestBackgroundWaitsForBusyGate(t *testing.T) {
|
||||
f := newFakeDixie(t, false)
|
||||
gate := NewGate(1, 0, 50*time.Millisecond) // no queue: a busy gate refuses at once
|
||||
e := f.engine(t, gate, true)
|
||||
doc := longDoc(100, 0)
|
||||
n := len(SplitPages(doc, ChunkTokens))
|
||||
if ok, _ := gate.Acquire(context.Background()); !ok {
|
||||
t.Fatal("could not take the gate")
|
||||
}
|
||||
// Run takes the gate before it chooses a retrieval mode, so queue the document directly.
|
||||
chunks := SplitPages(doc, ChunkTokens)
|
||||
e.queueForEmbed(chunks, "hash-g", "fleet")
|
||||
time.Sleep(2500 * time.Millisecond) // longer than the old single retry
|
||||
gate.Release()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
if err := e.WaitBackground(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := f.embedded.Load(); got != int64(n) {
|
||||
t.Fatalf("embedded %d of %d chunks: the worker gave up on a busy gate", got, n)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
+2
-1
@@ -61,7 +61,8 @@ PDFs work: `sift fetch` extracts them and returns the text. A short one comes ba
|
||||
long one comes back as a map — an outline with page numbers, then the first pages. Then:
|
||||
|
||||
- `sift fetch URL --pages 41-60` for a page range, `--grep 'pattern'` for matching lines, or
|
||||
`--question "…"` for a digest that cites each fact as `(p. N)`.
|
||||
`--question "…"` for a digest that cites each fact as `(p. N)`; `--question` works on long
|
||||
PDFs, and the first answer may be `bm25` (lexical) — ask again later for a sharper one.
|
||||
- Scanned PDFs have no text layer: sift reports `no text layer (scanned PDF?)` — point it at an
|
||||
OCR source rather than trying to fix it here.
|
||||
- A range returns at most 100 pages and grep at most 100 matches. On a very long document
|
||||
|
||||
Reference in New Issue
Block a user