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

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
2026-09-26 11:01:18 -07:00
co-authored by Claude Opus 5.5
parent 1311cf38d4
commit 49d96e197b
14 changed files with 1464 additions and 0 deletions
+63
View File
@@ -0,0 +1,63 @@
# P5b task 01: page-aware chunks and BM25
**Branch:** `p5b` (`git switch -c p5b master` if it does not exist, else `git switch p5b`; `git status --short` must be empty, otherwise stop)
**Commit subject:** `digest: SplitPages keeps every chunk on its page; BM25 scorer`
## Goal
Two pure functions for task 02. `SplitPages` is `Split` for text with `--- page N ---`
markers (what `pdfdoc.Full` produces): a chunk that starts inside a page begins with
`--- page N (cont.) ---`, so every chunk names its page. `BM25` ranks chunks lexically.
## Files
- Copy: `internal/digest/splitpages_test.go`, `internal/digest/bm25_test.go`
- Create: `internal/digest/splitpages.go`, `internal/digest/bm25.go`
- Modify: `docs/implementer-log.md` only
## Interfaces
```go
package digest
// SplitPages splits text into chunks of at most maxTok tokens (EstTokens) on line boundaries,
// like Split. A marker line matching ^--- page (\d+) ---$ sets the current page; a chunk that
// starts after the first marker but not at a marker line begins with
// "--- page N (cont.) ---\n\n" for the current page N. Text with no marker lines gives exactly
// Split's result.
func SplitPages(text string, maxTok int) []string
type BM25 struct{ /* unexported */ }
// NewBM25 indexes docs (k1 = 1.2, b = 0.75). Terms are the same tokens as terms() — lower-case
// runs of letters and digits, at least 2 characters — but counted, not as a set.
func NewBM25(docs []string) *BM25
// 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
```
## Rules
1. `SplitPages`: the continuation marker counts toward the chunk's budget (the test allows a
marker line of slack). Never split inside a line. Do not change `Split` (a P5a given test
pins it).
2. BM25 IDF: `ln(1 + (N − df + 0.5) / (df + 0.5))`, which is always positive. Average document
length over all documents; an all-empty corpus scores 0 everywhere.
## Steps
- [ ] **1.** Branch; copy the two given tests. **2. See them fail** (undefined names).
- [ ] **3.** `splitpages.go`, **4.** `bm25.go`.
- [ ] **5.** `gofmt -w`; `go test -race -count=3 ./internal/digest/` → `ok`. **6.** `make gate`.
- [ ] **7.** Row `p5b/01-splitpages-bm25`; commit.
```sh
git add internal/digest docs/implementer-log.md
git commit -F .state/commit-msg.txt
```
## Done when
- The given tests pass, every existing digest test passes, `make gate` ok, given files byte-identical.
+105
View File
@@ -0,0 +1,105 @@
# P5b task 02: cached embeddings and BM25-first retrieval for long documents
**Branch:** `p5b` (`git switch p5b`; `git status --short` must be empty, otherwise stop)
**Commit subject:** `digest: cache chunk embeddings; long documents answer from BM25 while embedding in the background`
## Goal
Today `selectText` embeds every chunk of a document for every question, inline, inside the
digest's 45 s timeout, and never uses the embedding cache. After this task:
- chunk embeddings are cached;
- a document with few uncached chunks is embedded inline and scored **hybrid** (as today);
- a long one is answered at once from **BM25**, with **no embedding call on the request path**,
while its chunks are embedded by a background worker;
- once they are, the next question is hybrid.
## Files
- Copy: `internal/digest/retrieval_test.go`
- Create: `internal/digest/retrieval.go` (selection), `internal/digest/background.go` (the worker)
- Modify: `internal/digest/digest.go` (`Result` fields, `Run`, `selectText`'s call site),
`docs/implementer-log.md`
## Interfaces
```go
package digest
const (
EmbedSyncMax = 48 // uncached chunks a question may embed inline; more go to the background
MaxDocChunks = 5000 // background embedding cap per document
)
// in Result:
Retrieval string // "whole", "head", "hybrid", "bm25" or "cached" (see rule 2)
Embedded int // chunks that had an embedding when the answer was built
// WaitBackground returns when the background worker is idle with nothing queued, or ctx.Err().
func (e *Engine) WaitBackground(ctx context.Context) error
```
What already exists and stays:
- `Split`, `SplitPages`, `NewBM25` (task 01);
- `EstTokens`, `SmallDoc`, `ChunkTokens`, `EmbedBatch`, `lexicalOverlap`, `cosine`;
- `e.embed(ctx, inputs)` (batches of `EmbedBatch`, halves on "too large");
- `e.Gate.Acquire/Release`;
- `dcache.EmbedKey(contentHash string, chunkIndex int, embedModel, callerClass string)`, `e.Cache.GetEmbed(key)` / `PutEmbed(key, vec)` (a nil `*DB` is a no-op).
## Rules
1. **Chunks and keys.**
- Chunk with `SplitPages(cleaned, ChunkTokens)`.
- Embedding keys are `dcache.EmbedKey(contentHash+"/sp1200", i, e.embedModel(), callerClass)`: the suffix binds the key to this chunker, so a future chunking change cannot reuse wrong vectors.
- `contentHash` and `callerClass` come from `Run` (pass them to the selection).
2. **Selection** (`retrieval.go`), in order:
- `EstTokens(cleaned) <= SmallDoc` → the whole text, `Retrieval "whole"`.
- Empty question → the leading chunks within `SmallDoc`, as today, `Retrieval "head"`.
- **No cache** (`e.Cache == nil`):
- `len(chunks) <= EmbedSyncMax` → embed everything inline, as today, `"hybrid"`;
- otherwise `"bm25"`, with no embedding and nothing queued (there is nowhere to keep the vectors).
- **Cache:** look up every chunk's vector.
- **uncached `<= EmbedSyncMax`** → embed the question and the uncached chunks inline (the question first, then the missing chunks, through `e.embed`), store each new chunk vector, `"hybrid"`.
- **otherwise** → `"bm25"`: rank by `NewBM25(chunks).Scores(question)` normalised by its maximum. Make **no embedding call**, and queue the document for the background (rule 4).
- Hybrid score = `0.7·cosine(q, v_i) + 0.3·lexicalOverlap(question, chunk)`, as today.
- Both modes then use today's greedy budget: chunk 0 first, then best-scoring chunks while the total stays within `SmallDoc`, emitted in document order.
- `Embedded` = the chunks with a vector at that moment.
3. **`Run`:**
- Set `r.Retrieval` and `r.Embedded`.
- On a digest-cache hit: `Retrieval "cached"`.
- **A digest made in `"bm25"` mode is not stored in the digest cache.** It is provisional; once the embeddings exist the same question must get a hybrid answer.
4. **Background worker** (`background.go`):
- One goroutine per `Engine`, started on first use (`sync.Once`), with a queue of at most **4 documents**, deduplicated by `contentHash+"|"+callerClass`. A full queue drops the document, logs `Warn` (with no text) and a later question re-queues it.
- For each document: embed the first `MaxDocChunks` chunks that lack a vector, **`EmbedBatch` at a time**. For each batch:
- `e.Gate.Acquire(context.Background())`, so interactive digests get turns; on a refusal (queue full or wait), sleep 1 s and try again;
- `e.embed` with a `e.timeout()` context;
- `Release`;
- store the vectors.
- On an embedding error: log `Warn`, drop the rest of that document, go on to the next.
- `WaitBackground` is true-idle: the queue is empty and no document is in progress. Poll every 20 ms or use a condition variable; either is fine.
5. Never log document text. The worker never holds the gate between batches.
## Facts
- The given test's fake dixie counts every embedded input. It can **hold `/embed` shut**, which is how it proves the `"bm25"` answer made no embedding call on the request path: if yours does, the first `Run` blocks and the test fails on time.
- `sseHandler` and `chatReq` come from `digest_test.go` / `digest.go` in the same package.
- Existing tests (`TestChunkOrderAndBudget`, `TestEmbedTooLargeRetriesAndCompletes`, …) run with no cache and must stay green.
## Steps
- [ ] **1.** `git switch p5b`; copy the given test. **2. See it fail.**
- [ ] **3.** `retrieval.go`, **4.** `background.go`, **5.** the `Run`/`Result` changes.
- [ ] **6.** `gofmt -w`; `go test -race -count=3 ./internal/digest/` → `ok`. **7.** `make gate`.
- [ ] **8.** Row `p5b/02-retrieval`; commit.
```sh
git add internal/digest docs/implementer-log.md
git commit -F .state/commit-msg.txt
```
## Done when
- The given and existing digest tests pass under `-race -count=3`; `make gate` ok; given file byte-identical.
## Stop and report if
- A given test's timing cannot be met without sleeping on the request path.
+43
View File
@@ -0,0 +1,43 @@
# P5b task 03: the API answers long-document questions and says how
**Branch:** `p5b` (`git switch p5b`; `git status --short` must be empty, otherwise stop)
**Commit subject:** `api: questions on long PDFs are answered (retrieval in the digest meta); README and skill`
## Goal
With task 02, the digest handles long documents itself. The API's 200-chunk refusal becomes a
5000-chunk safety cap (task 02's `MaxDocChunks`), and the `/v1/fetch` envelope's `digest`
object says how the answer was built.
## Files
- Copy: `internal/api/pdf_retrieval_test.go`, and the **replacement** `internal/api/pdf_test.go`
(P5a's copy asserted `DefaultPDFMaxQuestionChunks == 200`; now 5000 — the only change; the
P5b copy is the protected one from now on)
- Modify: `internal/api/pdf.go` (the constant), `internal/api/api.go` (the two digest-meta keys
only), `README.md`, `skill/SKILL.md`, `docs/implementer-log.md`
## Rules
1. `DefaultPDFMaxQuestionChunks = 5000`. The 422 message is unchanged.
2. In `v1fetch`, when the digest ran without a fallback, `digMeta` gains
`"retrieval": dr.Retrieval` and `"embedded": dr.Embedded`, next to `chunks_used` and
`chunks_total`.
3. **`README.md`:** in the PDF section, a question on a long document is answered first from BM25 (`retrieval: "bm25"`) while its chunks are embedded in the background, and is hybrid on later questions. Also mention the 5000-chunk cap.
4. **`skill/SKILL.md`:** one sentence: `--question` works on long PDFs; the first answer may be `bm25` (lexical); ask again later for a sharper one.
## Steps
- [ ] **1.** `git switch p5b`; copy the two given files. **2. See them fail.**
- [ ] **3.** The constant, the two meta keys, README, skill.
- [ ] **4.** `gofmt -w`; `go test -race -count=3 ./internal/api/` → `ok`. **5.** `make gate`.
- [ ] **6.** Row `p5b/03-api`; commit.
```sh
git add internal/api README.md skill/SKILL.md docs/implementer-log.md
git commit -F .state/commit-msg.txt
```
## Done when
- All tests pass; `make gate` ok; given files byte-identical.
+27
View File
@@ -0,0 +1,27 @@
# P5b implementation plan: questions on long documents
> **For the implementing model:** do not work from this file. The owner gives you one task file
> at a time. This file is the index for the owner and the reviewer.
**Goal:** P5a refuses `question` on a PDF longer than 200 digest chunks (~240k tokens) with
"use grep or pages". After P5b such a question is answered: at once from **BM25** over
page-aware chunks, while the document's chunks are **embedded in the background**; once they
are, later questions use today's hybrid score. Two facts found while planning:
- The embedding cache (`dcache` `GetEmbed`/`PutEmbed`, `EmbedKey`) exists but **nothing uses it**:
every question re-embeds every chunk.
- `Split` is line-based, so a chunk that starts mid-page carries no `--- page N ---` marker and
the model cannot cite its page.
| Task | Package | Given |
|---|---|---|
| 01-splitpages-bm25 | `internal/digest` (pure functions) | `digest/splitpages_test.go`, `digest/bm25_test.go` |
| 02-retrieval | `internal/digest` | `digest/retrieval_test.go` |
| 03-api | `internal/api`, `README.md`, `skill/SKILL.md` | `api/pdf_retrieval_test.go`; replaces `api/pdf_test.go` (P5a: the 200-chunk default becomes 5000) |
**How this plan was made:** acceptance tests first, no reference implementation; the given tests
compiled against a panic-only skeleton on master `1311cf3` and failed for the intended reasons.
## Changes during the run
(none yet)
@@ -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)
}
}
@@ -0,0 +1,403 @@
package api
// P5a task 06: PDFs through /v1/fetch and /v2/scrape — extraction in the child, the pool, the
// document cache, the map, pages, grep, and the errors.
//
// The child is the test binary itself (TestMain below), as in internal/pdftext's tests.
// SIFT_PDFTEXT_RUNLOG makes every child start append a line to a file, so a test can count
// extractions; SIFT_PDFTEXT_SLEEP slows a child down. Both hooks live only in this file.
import (
"encoding/json"
"fmt"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"net/netip"
"os"
"path/filepath"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"git.wntrmute.dev/kyle/sift/internal/auth"
"git.wntrmute.dev/kyle/sift/internal/dcache"
"git.wntrmute.dev/kyle/sift/internal/digest"
"git.wntrmute.dev/kyle/sift/internal/fetch"
"git.wntrmute.dev/kyle/sift/internal/guard"
"git.wntrmute.dev/kyle/sift/internal/pdfpool"
"git.wntrmute.dev/kyle/sift/internal/pdftest"
"git.wntrmute.dev/kyle/sift/internal/pdftext"
)
func TestMain(m *testing.M) {
if os.Getenv("SIFT_PDFTEXT_CHILD") == "1" {
if p := os.Getenv("SIFT_PDFTEXT_RUNLOG"); p != "" {
f, err := os.OpenFile(p, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600)
if err == nil {
_, _ = f.WriteString("run\n")
_ = f.Close()
}
}
if d, err := time.ParseDuration(os.Getenv("SIFT_PDFTEXT_SLEEP")); err == nil {
time.Sleep(d)
}
os.Exit(pdftext.ChildMain(os.Args[1:], os.Stdin, os.Stdout))
}
os.Exit(m.Run())
}
func threePDF() []byte {
return pdftest.Build(pdftest.Doc{
Title: "Three Pages",
Pages: []string{"Chapter one begins here.", "The quick brown fox is on page two.", "Last page."},
Outline: []pdftest.Outline{{Title: "1 Intro", Page: 1, Level: 1}, {Title: "2 Fox", Page: 2, Level: 1}},
})
}
func longPDF(n int) []byte {
var pages []string
for p := 1; p <= n; p++ {
pages = append(pages, pdftest.Lorem(p, 40))
}
return pdftest.Build(pdftest.Doc{Title: "Long Thesis", Pages: pages, Outline: []pdftest.Outline{{Title: "1 Introduction", Page: 1, Level: 1}, {Title: "2 Method", Page: n / 2, Level: 1}}})
}
// origin serves one PDF and counts GETs. It honours If-None-Match against its current ETag.
type origin struct {
mu sync.Mutex
ct string
body []byte
etag string
gets atomic.Int32
srv *httptest.Server
}
func newOrigin(t *testing.T, ct string, body []byte) *origin {
o := &origin{ct: ct, body: body, etag: `"v1"`}
o.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
o.gets.Add(1)
o.mu.Lock()
defer o.mu.Unlock()
if r.Header.Get("If-None-Match") == o.etag {
w.WriteHeader(http.StatusNotModified)
return
}
w.Header().Set("ETag", o.etag)
if o.ct != "" {
w.Header().Set("Content-Type", o.ct)
}
_, _ = w.Write(o.body)
}))
t.Cleanup(o.srv.Close)
return o
}
func (o *origin) replace(body []byte, etag string) {
o.mu.Lock()
o.body, o.etag = body, etag
o.mu.Unlock()
}
type pdfRig struct {
s *Server
fleet string // bearer key of caller "test" (hermes-fleet)
other string // bearer key of caller "other" (hermes-fleet)
gs string // bearer key of restricted caller "straylight-gs"
runlog string
logs *recHandler
}
func newPDFRig(t *testing.T, childEnv ...string) *pdfRig {
t.Helper()
dir := t.TempDir()
keys := filepath.Join(dir, "keys.toml")
key := func(name, ws string) string {
out := filepath.Join(dir, name)
if err := auth.Add(keys, name, ws, out); err != nil {
t.Fatal(err)
}
b, err := os.ReadFile(out)
if err != nil {
t.Fatal(err)
}
return strings.TrimSpace(string(b))
}
r := &pdfRig{fleet: key("test", "hermes-fleet"), other: key("other", "hermes-fleet"), gs: key("straylight-gs", "guardian-star"), runlog: filepath.Join(dir, "runs"), logs: &recHandler{}}
store, err := auth.Load(keys)
if err != nil {
t.Fatal(err)
}
g := &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}}
s := New("test", g, &fetch.Fetcher{Guard: g, UA: "sift-test"}, store, slog.New(r.logs))
s.PermitRestricted = []string{"straylight-gs"}
db, err := dcache.Open(filepath.Join(dir, "cache.db"), 0, nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = db.Close() })
s.Digest = &digest.Engine{Cache: db} // the cache lives on the digest engine; no model is configured
s.PDFRunner = &pdftext.Runner{Exe: os.Args[0], Env: append([]string{"SIFT_PDFTEXT_CHILD=1", "SIFT_PDFTEXT_RUNLOG=" + r.runlog}, childEnv...), Deadline: 20 * time.Second}
r.s = s
return r
}
func (r *pdfRig) runs() int {
b, _ := os.ReadFile(r.runlog)
return strings.Count(string(b), "run\n")
}
func (r *pdfRig) fetch(t *testing.T, key string, body map[string]any) (int, map[string]any, http.Header) {
t.Helper()
rec := postJSON(t, r.s, "/v1/fetch", key, body)
var out map[string]any
_ = json.Unmarshal(rec.Body.Bytes(), &out)
return rec.Code, out, rec.Header()
}
func pdfInfo(t *testing.T, out map[string]any) map[string]any {
t.Helper()
p, ok := out["pdf"].(map[string]any)
if !ok {
t.Fatalf("envelope has no pdf object: %v", out)
}
return p
}
func TestPDFCleanFetch(t *testing.T) {
for _, ct := range []string{"application/pdf", "application/octet-stream", ""} {
r := newPDFRig(t)
o := newOrigin(t, ct, threePDF())
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
if code != 200 {
t.Fatalf("%q: %d %v", ct, code, out)
}
content, _ := out["content"].(string)
if !strings.Contains(content, "--- page 2 ---") || !strings.Contains(content, "quick brown fox") || out["mode"] != "clean" {
t.Errorf("%q: content %q mode %v", ct, content, out["mode"])
}
p := pdfInfo(t, out)
if p["pages"] != 3.0 || p["title"] != "Three Pages" || p["map"] != false || p["first"] != 1.0 || p["last"] != 3.0 {
t.Errorf("%q: pdf = %v", ct, p)
}
if ol, _ := p["outline"].([]any); len(ol) != 2 {
t.Errorf("%q: outline = %v", ct, p["outline"])
}
}
}
func TestPDFMap(t *testing.T) {
r := newPDFRig(t)
r.s.OutCap = 20000
o := newOrigin(t, "application/pdf", longPDF(40)) // ~40 × 2,700 characters
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
if code != 200 {
t.Fatalf("%d %v", code, out)
}
content, _ := out["content"].(string)
p := pdfInfo(t, out)
if p["map"] != true || !strings.Contains(content, "## Outline") || !strings.Contains(content, "- 2 Method (p. 20)") || len([]rune(content)) > 20000 {
t.Errorf("map: pdf %v, %d chars:\n%.400s", p, len([]rune(content)), content)
}
}
func TestPDFPagesAndGrep(t *testing.T) {
r := newPDFRig(t)
o := newOrigin(t, "application/pdf", threePDF())
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "pages": "2-3"})
content, _ := out["content"].(string)
if code != 200 || strings.Contains(content, "--- page 1 ---") || !strings.Contains(content, "--- page 3 ---") {
t.Errorf("pages 2-3: %d %q", code, content)
}
if p := pdfInfo(t, out); p["first"] != 2.0 || p["last"] != 3.0 {
t.Errorf("pages 2-3: pdf %v", p)
}
for _, bad := range []map[string]any{{"pages": "0-1"}, {"pages": "9"}, {"grep": strings.Repeat("a", 201)}, {"grep": "(unclosed"}} {
bad["url"] = o.srv.URL
if code, out, _ := r.fetch(t, r.fleet, bad); code != 400 || out["success"] != false {
t.Errorf("%v: %d %v, want 400", bad, code, out)
}
}
code, out, _ = r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "grep": "quick brown"})
content, _ = out["content"].(string)
if code != 200 || !strings.Contains(content, "p. 2: The quick brown fox is on page two.") || pdfInfo(t, out)["grep_hits"] != 1.0 {
t.Errorf("grep: %d %q %v", code, content, out["pdf"])
}
// pages and grep on a non-PDF are refused rather than ignored.
html := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "text/html")
_, _ = io.WriteString(w, "<html><body><p>hi</p></body></html>")
}))
defer html.Close()
if code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": html.URL, "pages": "1"}); code != 400 {
t.Errorf("pages on html: %d %v, want 400", code, out)
}
}
func TestPDFCache(t *testing.T) {
r := newPDFRig(t)
o := newOrigin(t, "application/pdf", threePDF())
get := func(extra map[string]any) string {
body := map[string]any{"url": o.srv.URL}
for k, v := range extra {
body[k] = v
}
code, out, _ := r.fetch(t, r.fleet, body)
if code != 200 {
t.Fatalf("%v: %d %v", extra, code, out)
}
c, _ := out["content"].(string)
return c
}
get(nil)
if o.gets.Load() != 1 || r.runs() != 1 {
t.Fatalf("first fetch: %d GETs, %d extractions; want 1 and 1", o.gets.Load(), r.runs())
}
if c := get(map[string]any{"pages": "2"}); !strings.Contains(c, "quick brown fox") {
t.Errorf("cached pages=2: %q", c)
}
if o.gets.Load() != 1 || r.runs() != 1 {
t.Errorf("fresh cache hit: %d GETs, %d extractions; want still 1 and 1", o.gets.Load(), r.runs())
}
// Stale: a conditional GET, answered 304, re-uses the text.
r.s.PDFTTL = time.Nanosecond
time.Sleep(2 * time.Millisecond)
get(map[string]any{"grep": "Chapter"})
if o.gets.Load() != 2 || r.runs() != 1 {
t.Errorf("after 304: %d GETs, %d extractions; want 2 and 1", o.gets.Load(), r.runs())
}
// The document changed: a full GET and a new extraction.
o.replace(pdftest.Build(pdftest.Doc{Pages: []string{"Brand new text."}}), `"v2"`)
if c := get(nil); !strings.Contains(c, "Brand new text.") {
t.Errorf("changed document: %q", c)
}
if o.gets.Load() != 3 || r.runs() != 2 {
t.Errorf("after change: %d GETs, %d extractions; want 3 and 2", o.gets.Load(), r.runs())
}
if DefaultPDFTTL != 24*time.Hour {
t.Errorf("DefaultPDFTTL = %v, want 24h", DefaultPDFTTL)
}
}
func TestPDFErrors(t *testing.T) {
r := newPDFRig(t)
enc, err := os.ReadFile(filepath.Join("..", "..", "tests", "testdata", "pdf", "encrypted-user.pdf"))
if err != nil {
t.Fatal(err)
}
for _, tc := range []struct {
name string
body []byte
want string
hint string
}{
{"scanned", pdftest.Build(pdftest.Doc{Pages: []string{"", ""}, Scanned: []int{1, 2}}), "no text layer", "ocr"},
{"encrypted", enc, "encrypted", ""},
{"garbage", []byte("%PDF-1.4\nthis is not really a pdf"), "unreadable", ""},
} {
o := newOrigin(t, "application/pdf", tc.body)
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
msg, _ := out["error"].(string)
if code != 422 || !strings.Contains(msg, tc.want) || (tc.hint != "" && out["hint"] != tc.hint) {
t.Errorf("%s: %d %v; want 422 %q hint %q", tc.name, code, out, tc.want, tc.hint)
}
}
}
func TestPDFTimeout(t *testing.T) {
r := newPDFRig(t, "SIFT_PDFTEXT_SLEEP=10s")
r.s.PDFRunner.Deadline = 300 * time.Millisecond
o := newOrigin(t, "application/pdf", threePDF())
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL})
if msg, _ := out["error"].(string); code != 422 || !strings.Contains(msg, "timed out") {
t.Errorf("slow child: %d %v; want 422 … timed out", code, out)
}
}
func TestPDFBusy(t *testing.T) {
r := newPDFRig(t, "SIFT_PDFTEXT_SLEEP=1500ms")
r.s.PDFPool = pdfpool.New(1, 0, 100*time.Millisecond)
a := newOrigin(t, "application/pdf", threePDF())
b := newOrigin(t, "application/pdf", longPDF(2))
done := make(chan int, 1)
go func() {
code, _, _ := r.fetch(t, r.fleet, map[string]any{"url": a.srv.URL})
done <- code
}()
time.Sleep(400 * time.Millisecond) // the first extraction holds the only slot
code, out, hdr := r.fetch(t, r.other, map[string]any{"url": b.srv.URL})
if code != 503 || hdr.Get("Retry-After") != "10" || !strings.Contains(fmt.Sprint(out["error"]), "busy") {
t.Errorf("second caller with the pool full: %d %v Retry-After %q; want 503, busy, 10", code, out, hdr.Get("Retry-After"))
}
code, out, _ = r.fetch(t, r.fleet, map[string]any{"url": b.srv.URL})
if code != 429 {
t.Errorf("same caller, second document while its first is running: %d %v; want 429", code, out)
}
if c := <-done; c != 200 {
t.Errorf("first fetch: %d", c)
}
}
func TestPDFRestrictedCaller(t *testing.T) {
r := newPDFRig(t)
o := newOrigin(t, "application/pdf", threePDF())
if code, out, _ := r.fetch(t, r.gs, map[string]any{"url": o.srv.URL}); code != 200 {
t.Fatalf("restricted fetch: %d %v", code, out)
}
r.logs.mu.Lock()
for _, rec := range r.logs.recs {
var b strings.Builder
b.WriteString(rec.Message)
rec.Attrs(func(a slog.Attr) bool { fmt.Fprintf(&b, " %s=%v", a.Key, a.Value); return true })
line := b.String()
if strings.Contains(line, "127.0.0.1") || strings.Contains(line, "Three Pages") || strings.Contains(line, "quick brown") {
t.Errorf("a restricted caller's PDF fetch logged %q", line)
}
}
r.logs.mu.Unlock()
// The fleet does not see the restricted caller's cached text: it extracts its own copy.
if code, _, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL}); code != 200 || r.runs() != 2 {
t.Errorf("fleet after restricted: %d, %d extractions; want 200 and 2 (separate cache partitions)", code, r.runs())
}
}
func TestPDFScrape(t *testing.T) {
r := newPDFRig(t)
o := newOrigin(t, "application/pdf", threePDF())
rec := postJSON(t, r.s, "/v2/scrape", r.fleet, map[string]any{"url": o.srv.URL, "formats": []string{"markdown"}})
var out struct {
Success bool `json:"success"`
Data struct {
Markdown string `json:"markdown"`
Metadata map[string]any `json:"metadata"`
} `json:"data"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil || rec.Code != 200 || !out.Success {
t.Fatalf("%d %s", rec.Code, rec.Body.String())
}
if !strings.Contains(out.Data.Markdown, "--- page 2 ---") || out.Data.Metadata["numPages"] != 3.0 || out.Data.Metadata["title"] != "Three Pages" {
t.Errorf("scrape: markdown %q metadata %v", out.Data.Markdown, out.Data.Metadata)
}
if ct, _ := out.Data.Metadata["contentType"].(string); !strings.HasPrefix(ct, "application/pdf") {
t.Errorf("contentType = %q", ct)
}
}
// A question on a document longer than PDFMaxQuestionChunks is refused with a pointer to grep
// and pages, before any model call.
func TestPDFQuestionTooLong(t *testing.T) {
r := newPDFRig(t)
r.s.PDFMaxQuestionChunks = 2
o := newOrigin(t, "application/pdf", longPDF(12))
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "question": "what is the method?"})
if msg, _ := out["error"].(string); code != 422 || !strings.Contains(msg, "too long") || !strings.Contains(msg, "grep") {
t.Errorf("long question: %d %v; want 422 … too long … grep", code, out)
}
// P5b: long documents are answered by retrieval now; the limit is only a safety cap.
if DefaultPDFMaxQuestionChunks != 5000 {
t.Errorf("DefaultPDFMaxQuestionChunks = %d, want 5000", DefaultPDFMaxQuestionChunks)
}
}
@@ -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")
}
}
}
@@ -0,0 +1,244 @@
package digest
// P5b task 02: retrieval for long documents. Embeddings are cached per chunk; a document with
// more uncached chunks than EmbedSyncMax is answered at once from BM25 while its chunks are
// embedded in the background, and later questions use the hybrid score.
import (
"context"
"encoding/json"
"fmt"
"hash/fnv"
"net/http"
"net/http/httptest"
"path/filepath"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"git.wntrmute.dev/kyle/sift/internal/dcache"
)
// fakeDixie serves /embed (a bag-of-words hash into 64 dims, so cosine similarity means
// something) and /chat (the SSE stub). It counts embedded inputs and can hold /embed shut.
type fakeDixie struct {
srv *httptest.Server
embedded atomic.Int64 // inputs embedded, questions included
hold chan struct{}
holdOnce sync.Once
lastChat atomic.Pointer[chatReq]
}
func newFakeDixie(t *testing.T, hold bool) *fakeDixie {
f := &fakeDixie{}
if hold {
f.hold = make(chan struct{})
}
var disc atomic.Bool
mux := http.NewServeMux()
mux.HandleFunc("/embed", func(w http.ResponseWriter, r *http.Request) {
if f.hold != nil {
select {
case <-f.hold:
case <-r.Context().Done():
return
}
}
var body struct {
Input []string `json:"input"`
}
_ = json.NewDecoder(r.Body).Decode(&body)
var data []map[string]any
for i, s := range body.Input {
data = append(data, map[string]any{"index": i, "embedding": bow(s)})
}
f.embedded.Add(int64(len(body.Input)))
_ = json.NewEncoder(w).Encode(map[string]any{"data": data})
})
mux.Handle("/chat", sseHandler(t, "ok", &f.lastChat, &disc))
f.srv = httptest.NewServer(mux)
t.Cleanup(func() { f.release(); f.srv.Close() })
return f
}
func (f *fakeDixie) release() {
if f.hold != nil {
f.holdOnce.Do(func() { close(f.hold) })
}
}
func bow(s string) []float64 {
v := make([]float64, 64)
for t := range terms(s) {
h := fnv.New32a()
_, _ = h.Write([]byte(t))
v[h.Sum32()%64]++
}
return v
}
func (f *fakeDixie) engine(t *testing.T, gate *Gate, withCache bool) *Engine {
e := &Engine{ChatURL: f.srv.URL + "/chat", EmbedURL: f.srv.URL + "/embed", Gate: gate, Client: f.srv.Client()}
if withCache {
db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 0, nil)
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = db.Close() })
e.Cache = db
}
return e
}
// longDoc is n pages of distinct filler; the page named rare carries a rare term.
func longDoc(n, rare int) string {
var b strings.Builder
words := []string{"harbour", "vessel", "cargo", "anchor", "pilot", "tug", "berth", "quay", "hull", "keel", "mast", "deck"}
for p := 1; p <= n; p++ {
if p > 1 {
b.WriteString("\n\n")
}
fmt.Fprintf(&b, "--- page %d ---\n\n", p)
for l := 0; l < 40; l++ {
fmt.Fprintf(&b, "page %d line %d %s %s %s %s\n", p, l, words[(p+l)%12], words[(p*3+l)%12], words[(p*5+l)%12], words[(p*7+l)%12])
}
if p == rare {
b.WriteString("the zanzibarite cargo manifest was signed by the second officer\n")
}
}
return b.String()
}
func chatDoc(f *fakeDixie) string {
req := f.lastChat.Load()
if req == nil {
return ""
}
var s strings.Builder
for _, m := range req.Messages {
s.WriteString(m.Content)
}
return s.String()
}
func TestEmbeddingsAreCachedAcrossQuestions(t *testing.T) {
f := newFakeDixie(t, false)
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
doc := longDoc(12, 0) // ~12 pages, a few chunks: under EmbedSyncMax
r := e.Run(context.Background(), doc, "which berth did the tug use?", "https://ex.test/a.pdf", dcache.ClassFleet)
if r.Fallback != "" || r.Retrieval != "hybrid" {
t.Fatalf("first question: fallback %q retrieval %q", r.Fallback, r.Retrieval)
}
n := len(SplitPages(doc, ChunkTokens))
if got := f.embedded.Load(); got != int64(n+1) {
t.Fatalf("first question embedded %d inputs, want %d chunks + the question", got, n)
}
r = e.Run(context.Background(), doc, "where was the pilot?", "https://ex.test/a.pdf", dcache.ClassFleet)
if r.Fallback != "" || r.Retrieval != "hybrid" || r.Embedded != n {
t.Fatalf("second question: %+v", r)
}
if got := f.embedded.Load(); got != int64(n+2) {
t.Errorf("second question embedded %d inputs in total, want %d: only its own question is new", got, n+2)
}
}
// A long document is answered from BM25 without waiting for a single embedding, and becomes
// hybrid once the background has embedded every chunk.
func TestLongDocumentBM25ThenHybrid(t *testing.T) {
f := newFakeDixie(t, true) // /embed held shut until release
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
doc := longDoc(300, 250)
n := len(SplitPages(doc, ChunkTokens))
if n <= EmbedSyncMax {
t.Fatalf("test document has %d chunks; it must exceed EmbedSyncMax (%d)", n, EmbedSyncMax)
}
start := time.Now()
r := e.Run(context.Background(), doc, "who signed the zanzibarite manifest?", "https://ex.test/big.pdf", dcache.ClassFleet)
if d := time.Since(start); d > 10*time.Second {
t.Fatalf("first answer took %v: it waited on embeddings", d)
}
if r.Fallback != "" || r.Retrieval != "bm25" || r.Embedded != 0 {
t.Fatalf("first answer: %+v, want retrieval bm25 with 0 embedded", r)
}
if doc := chatDoc(f); !strings.Contains(doc, "zanzibarite") || !strings.Contains(doc, "page 250") {
t.Errorf("BM25 did not bring page 250 (the only one with the term) into the prompt")
}
f.release()
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
if err := e.WaitBackground(ctx); err != nil {
t.Fatalf("background embedding did not finish: %v", err)
}
r = e.Run(context.Background(), doc, "who signed the zanzibarite manifest?", "https://ex.test/big.pdf", dcache.ClassFleet)
if r.Fallback != "" || r.Retrieval != "hybrid" || r.Embedded != n {
t.Fatalf("after background: %+v, want hybrid with all %d embedded", r, n)
}
if doc := chatDoc(f); !strings.Contains(doc, "zanzibarite") {
t.Errorf("hybrid lost the page with the term")
}
}
// Two questions on the same long document queue it once: every chunk is embedded once.
func TestBackgroundDedupes(t *testing.T) {
f := newFakeDixie(t, true)
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
doc := longDoc(200, 0)
n := len(SplitPages(doc, ChunkTokens))
e.Run(context.Background(), doc, "first question about the quay", "https://ex.test/b.pdf", dcache.ClassFleet)
e.Run(context.Background(), doc, "second question about the keel", "https://ex.test/b.pdf", dcache.ClassFleet)
f.release()
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
if err := e.WaitBackground(ctx); err != nil {
t.Fatal(err)
}
if got := f.embedded.Load(); got > int64(n) {
t.Errorf("embedded %d inputs for a %d-chunk document: queued twice", got, n)
}
}
// Background work takes the gate one batch at a time, so an interactive digest is not starved.
func TestBackgroundYieldsTheGate(t *testing.T) {
f := newFakeDixie(t, false)
e := f.engine(t, NewGate(1, 8, 10*time.Second), true)
e.Run(context.Background(), longDoc(400, 0), "start the background", "https://ex.test/c.pdf", dcache.ClassFleet)
start := time.Now()
r := e.Run(context.Background(), "a short page about a tug and a quay", "what is here?", "https://ex.test/small", dcache.ClassFleet)
if r.Fallback != "" {
t.Fatalf("interactive digest fell back (%q) while background embedding ran", r.Fallback)
}
if d := time.Since(start); d > 5*time.Second {
t.Errorf("interactive digest waited %v behind background embedding", d)
}
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
_ = e.WaitBackground(ctx)
}
// Without a cache there is nowhere to keep background work: a long document is BM25 and
// nothing is embedded.
func TestNoCacheLongDocumentIsBM25Only(t *testing.T) {
f := newFakeDixie(t, false)
e := f.engine(t, NewGate(2, 8, 5*time.Second), false)
r := e.Run(context.Background(), longDoc(300, 250), "who signed the zanzibarite manifest?", "https://ex.test/d.pdf", dcache.ClassFleet)
if r.Fallback != "" || r.Retrieval != "bm25" {
t.Fatalf("%+v, want bm25", r)
}
time.Sleep(200 * time.Millisecond)
if got := f.embedded.Load(); got != 0 {
t.Errorf("embedded %d inputs with no cache to keep them in", got)
}
}
// A document that fits the digest budget whole is sent whole, as today.
func TestSmallDocumentIsWhole(t *testing.T) {
f := newFakeDixie(t, false)
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
r := e.Run(context.Background(), "--- page 1 ---\n\nOne short page.", "what is it?", "https://ex.test/e.pdf", dcache.ClassFleet)
if r.Fallback != "" || r.Retrieval != "whole" {
t.Fatalf("%+v, want whole", r)
}
}
@@ -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))
}
}
+76
View File
@@ -0,0 +1,76 @@
# P5c task 01: `POST /v1/extract`
**Branch:** `p5c` (`git switch -c p5c master` if it does not exist, else `git switch p5c`; `git status --short` must be empty, otherwise stop)
**Commit subject:** `api: POST /v1/extract answers an uploaded PDF like a fetched one (content-addressed cache)`
## Goal
`POST /v1/extract` with a PDF as the request body returns the `/v1/fetch` PDF envelope:
text with page markers or the map, `pages`, `grep`, or a cited digest. Nothing is fetched.
The bytes are hashed, the cache key is the hash, and the same file is extracted once.
## Files
- Copy: `internal/api/extract_test.go`
- Create: `internal/api/extract.go` (the handler and its envelope)
- Modify: `internal/api/api.go`: only the route in `New`, the `Server` field and constant
below, and the two lines in `ServeHTTP` (rule 1); `docs/implementer-log.md`
## Interfaces
```go
package api
const DefaultPDFMaxUpload = 100 << 20
// in Server:
PDFMaxUpload int64 // 0 = DefaultPDFMaxUpload
// route: s.Mux.HandleFunc("POST /v1/extract", s.v1extract)
// query parameters (all optional): pages, grep, question, mode ("clean" | "digest";
// default digest when question is set, else clean)
```
What exists and is reused, unchanged:
- `s.caller`, `s.callerClass`;
- `s.extractPDF(ctx, c, *fetch.Resource)`, `pdfError`, `s.renderPDF`, `s.storePDF(key, res, doc)`, `s.pdfFromEntry`;
- `s.cache()`, `dcache.PDFKey`, `clean.Kind`, `pdfdoc.Full`;
- `s.Digest.Run`, `s.Prom.PDFCacheHit`, `writeErr`, `writeJSON`.
## Rules
1. **`ServeHTTP`:**
- today it wraps every request body in `http.MaxBytesReader(w, r.Body, maxReqBody)` (64 KiB). For `/v1/extract` use `PDFMaxUpload` instead;
- add `/v1/extract` to the paths that get `requestDeadline`;
- every other path keeps 64 KiB (the given test checks that).
2. **Handler:**
- Authenticate as the other endpoints do.
- Read the whole body. Empty → `400 empty body`. Over the limit (`*http.MaxBytesError`) → `413 too large (>N MB)`.
- `clean.Kind(Content-Type, body)` must be `"pdf"`; otherwise `415 extract accepts PDFs only`.
3. **Cache:**
- `hash := "sha256:" + hex(sha256(body))`, `key := dcache.PDFKey("upload:"+hash, s.callerClass(c))`.
- An entry → the doc from it, `cached: true`, `PDFCacheHit`, no child. The content is addressed by its hash, so there is **no TTL and no revalidation**.
- Otherwise build `&fetch.Resource{URL: "upload", FinalURL: "upload", Status: 200, ContentType: ct, Body: body, Hash: hash, FetchedAt: now}`, then `extractPDF`, then `storePDF(key, res, doc)`.
- Errors go through `pdfError`, via the existing `writeFetchErr` (it adds `Retry-After`).
4. **Answer:** the same fields as `/v1/fetch` for a PDF:
- `content`, `mode`, `untrusted: true`, `warnings`, `sizes`, `timings`, `pdf`, and `digest` when a digest ran, including P5b's `retrieval` and `embedded`;
- `source` is `{"kind": "upload", "hash", "bytes", "cached", "fetched_at"}`, with **no `url`** key;
- bad `pages`/`grep` → 400 (as in `v1fetch`); the question limit → the same 422 as `v1fetch`.
- Do not call Honcho context for uploads.
5. **Logging:** one line, `"extract"`, with caller, bytes, `pdf_pages`, mode and cached. Never a title or any text. The server never sees a filename.
## Steps
- [ ] **1.** Branch; copy the given test. **2. See it fail.**
- [ ] **3.** `api.go` hooks, **4.** `extract.go`.
- [ ] **5.** `gofmt -w`; `go test -race -count=3 ./internal/api/` → `ok`. **6.** `make gate`.
- [ ] **7.** Row `p5c/01-extract-api`; commit.
```sh
git add internal/api docs/implementer-log.md
git commit -F .state/commit-msg.txt
```
## Done when
- All tests pass under `-race -count=3`; `make gate` ok; given file byte-identical.
+60
View File
@@ -0,0 +1,60 @@
# P5c task 02: `sift extract <file.pdf>`
**Branch:** `p5c` (`git switch p5c`; `git status --short` must be empty, otherwise stop)
**Commit subject:** `cli: sift extract uploads a local PDF; README and skill`
## Goal
`sift extract thesis.pdf [--pages 41-60] [--grep RE] [--question Q] [--json]` uploads the file
to `/v1/extract` and prints the answer exactly as `sift fetch` does.
## Files
- Copy: `cmd/sift/extract_test.go`
- Create: `cmd/sift/extractcmd.go`
- Modify: `cmd/sift/main.go` (a `case "extract":` and one help line), `README.md`,
`skill/SKILL.md`, `docs/implementer-log.md`
## Interfaces
```go
package main
type extractOpts struct {
File, Pages, Grep, Question, Mode string
JSON bool
}
const extractUsage = "usage: sift extract <file.pdf> [--pages a-b] [--grep RE] [--question Q] [--mode clean|digest] [--json]"
func parseExtractOpts(args []string) (extractOpts, error) // flags before or after the file, as fetch
func extractURL(base string, o extractOpts) string // base + "/v1/extract?" + only the options that are set; mode=digest when Question is set and Mode is empty
```
## Rules
1. The command:
- Resolve the base URL and key exactly as `sift fetch` does (reuse its client-config helpers; do not copy them).
- `stat` the file: over 100 MiB → exit with a message before uploading.
- `POST` the bytes with `Content-Type: application/pdf` and the bearer key.
- Print `content` (or the whole envelope with `--json`), with the same errors and exit codes as fetch (`exitFor`).
- Do not read the whole file into memory twice: stream it with `os.Open` as the request body and set `ContentLength`.
2. **`README.md`:** a `sift extract` subsection next to `sift fetch`.
3. **`skill/SKILL.md`:** "A PDF on disk (an attachment, a download): `sift extract <file>`, with the same `--pages`/`--grep`/`--question`. Never run `pdftotext` yourself."
## Steps
- [ ] **1.** `git switch p5c`; copy the given test. **2. See it fail.**
- [ ] **3.** `extractcmd.go`, `main.go`, README, skill.
- [ ] **4.** `gofmt -w`; `go test -race -count=3 ./cmd/sift/` → `ok`. **5.** `make gate`.
- [ ] **6.** `make build && bin/sift extract 2>&1 | head -1` prints the usage line.
- [ ] **7.** Row `p5c/02-extract-cli`; commit.
```sh
git add cmd/sift README.md skill/SKILL.md docs/implementer-log.md
git commit -F .state/commit-msg.txt
```
## Done when
- All tests pass; `make gate` ok; given file byte-identical.
+21
View File
@@ -0,0 +1,21 @@
# P5c implementation plan: PDFs already on a host
> **For the implementing model:** do not work from this file. The owner gives you one task file
> at a time. This file is the index for the owner and the reviewer.
**Goal:** an agent with a PDF on disk (a Mattermost attachment, a download) gets the same
answer as for a PDF URL — text with page markers, the map, pages, grep, cited digests — by
uploading the bytes: `POST /v1/extract`, CLI `sift extract <file>`. No poppler on the host, no
shell `pdftotext`. The cache key is the content hash, so the same file is extracted once.
| Task | Package | Given |
|---|---|---|
| 01-extract-api | `internal/api` | `api/extract_test.go` |
| 02-extract-cli | `cmd/sift`, `README.md`, `skill/SKILL.md` | `cmd/sift/extract_test.go` |
Runs after P5b (task 01 uses the P5b digest fields). Given tests compiled against a panic-only
skeleton on master `1311cf3` and failed for the intended reasons.
## Changes during the run
(none yet)
@@ -0,0 +1,46 @@
package main
// P5c task 02: `sift extract <file>` uploads a PDF to /v1/extract.
import (
"net/url"
"strings"
"testing"
)
func TestParseExtractOpts(t *testing.T) {
o, err := parseExtractOpts([]string{"thesis.pdf", "--pages", "41-60", "--grep", "attention head"})
if err != nil || o.File != "thesis.pdf" || o.Pages != "41-60" || o.Grep != "attention head" {
t.Errorf("%+v %v", o, err)
}
o, err = parseExtractOpts([]string{"--question", "what is the method?", "--json", "paper.pdf"})
if err != nil || o.File != "paper.pdf" || o.Question != "what is the method?" || !o.JSON {
t.Errorf("flags first: %+v %v", o, err)
}
if _, err := parseExtractOpts(nil); err == nil {
t.Error("no file: want an error")
}
if _, err := parseExtractOpts([]string{"a.pdf", "b.pdf"}); err == nil {
t.Error("two files: want an error")
}
if !strings.Contains(extractUsage, "sift extract <file.pdf>") {
t.Errorf("usage %q", extractUsage)
}
}
func TestExtractURL(t *testing.T) {
u, err := url.Parse(extractURL("https://sift.example:8443", extractOpts{File: "x.pdf", Pages: "2-3", Grep: "a b&c"}))
if err != nil {
t.Fatal(err)
}
if u.Path != "/v1/extract" || u.Query().Get("pages") != "2-3" || u.Query().Get("grep") != "a b&c" {
t.Errorf("url %s", u)
}
if u.Query().Has("question") || u.Query().Has("mode") {
t.Errorf("unset options must not be sent: %s", u)
}
u, _ = url.Parse(extractURL("https://sift.example:8443", extractOpts{File: "x.pdf", Question: "why?"}))
if u.Query().Get("question") != "why?" || u.Query().Get("mode") != "digest" {
t.Errorf("question url %s (mode must be digest with a question)", u)
}
}
@@ -0,0 +1,122 @@
package api
// P5c task 01: POST /v1/extract takes a PDF's bytes (a file already on the caller's host), and
// answers exactly like /v1/fetch does for a PDF URL. The cache key is the content hash.
import (
"bytes"
"encoding/json"
"fmt"
"log/slog"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
func (r *pdfRig) extract(t *testing.T, key, query string, body []byte, ct string) (int, map[string]any) {
t.Helper()
req := httptest.NewRequest(http.MethodPost, "/v1/extract"+query, bytes.NewReader(body))
if ct != "" {
req.Header.Set("Content-Type", ct)
}
if key != "" {
req.Header.Set("Authorization", "Bearer "+key)
}
rec := httptest.NewRecorder()
r.s.ServeHTTP(rec, req)
var out map[string]any
_ = json.Unmarshal(rec.Body.Bytes(), &out)
return rec.Code, out
}
func TestExtractUpload(t *testing.T) {
r := newPDFRig(t)
pdf := threePDF()
code, out := r.extract(t, r.fleet, "", pdf, "application/pdf")
if code != 200 {
t.Fatalf("%d %v", code, out)
}
content, _ := out["content"].(string)
if !strings.Contains(content, "--- page 2 ---") || !strings.Contains(content, "quick brown fox") {
t.Errorf("content %q", content)
}
if p := pdfInfo(t, out); p["pages"] != 3.0 || p["title"] != "Three Pages" {
t.Errorf("pdf %v", p)
}
src, _ := out["source"].(map[string]any)
if src["kind"] != "upload" || src["bytes"] != float64(len(pdf)) || !strings.HasPrefix(fmt.Sprint(src["hash"]), "sha256:") || src["cached"] != false {
t.Errorf("source %v", src)
}
if _, hasURL := src["url"]; hasURL {
t.Error("an upload has no URL; source must not invent one")
}
// The same bytes again: from the cache, no child.
code, out = r.extract(t, r.fleet, "?pages=2", pdf, "application/octet-stream")
if code != 200 || r.runs() != 1 {
t.Fatalf("second upload: %d, %d extractions; want 200 and 1", code, r.runs())
}
if src, _ := out["source"].(map[string]any); src["cached"] != true {
t.Errorf("second upload source %v, want cached", src)
}
if c, _ := out["content"].(string); strings.Contains(c, "--- page 1 ---") || !strings.Contains(c, "--- page 2 ---") {
t.Errorf("pages=2: %q", c)
}
code, out = r.extract(t, r.fleet, "?grep=quick+brown", pdf, "application/pdf")
if c, _ := out["content"].(string); code != 200 || !strings.Contains(c, "p. 2: The quick brown fox") {
t.Errorf("grep: %d %q", code, c)
}
}
func TestExtractRefusals(t *testing.T) {
r := newPDFRig(t)
if code, _ := r.extract(t, "", "", threePDF(), "application/pdf"); code != 401 {
t.Errorf("no key: %d, want 401", code)
}
if code, out := r.extract(t, r.fleet, "", []byte("<html>not a pdf</html>"), "text/html"); code != 415 {
t.Errorf("html body: %d %v, want 415", code, out)
}
if code, _ := r.extract(t, r.fleet, "", nil, "application/pdf"); code != 400 {
t.Errorf("empty body: %d, want 400", code)
}
if code, _ := r.extract(t, r.fleet, "?pages=0-1", threePDF(), "application/pdf"); code != 400 {
t.Errorf("bad pages: %d, want 400", code)
}
r.s.PDFMaxUpload = 1000
if code, out := r.extract(t, r.fleet, "", threePDF(), "application/pdf"); code != 413 {
t.Errorf("over PDFMaxUpload: %d %v, want 413", code, out)
}
if DefaultPDFMaxUpload != 100<<20 {
t.Errorf("DefaultPDFMaxUpload = %d, want 100 MiB", DefaultPDFMaxUpload)
}
// The ordinary JSON endpoints keep their small request limit.
big := bytes.Repeat([]byte("x"), 128<<10)
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/v1/fetch", bytes.NewReader(append([]byte(`{"url":"`), append(big, []byte(`"}`)...)...)))
req.Header.Set("Authorization", "Bearer "+r.fleet)
r.s.ServeHTTP(rec, req)
if rec.Code == 200 {
t.Error("a 128 KiB /v1/fetch body was accepted; the JSON limit must still apply there")
}
}
func TestExtractRestrictedPartition(t *testing.T) {
r := newPDFRig(t)
pdf := threePDF()
if code, _ := r.extract(t, r.gs, "", pdf, "application/pdf"); code != 200 {
t.Fatalf("restricted upload: %d", code)
}
r.logs.mu.Lock()
for _, rec := range r.logs.recs {
var b strings.Builder
b.WriteString(rec.Message)
rec.Attrs(func(a slog.Attr) bool { fmt.Fprintf(&b, " %s=%v", a.Key, a.Value); return true })
if strings.Contains(b.String(), "Three Pages") || strings.Contains(b.String(), "quick brown") {
t.Errorf("an upload's log line carries document text: %q", b.String())
}
}
r.logs.mu.Unlock()
if code, _ := r.extract(t, r.fleet, "", pdf, "application/pdf"); code != 200 || r.runs() != 2 {
t.Errorf("fleet after restricted: %d, %d extractions; want 200 and 2 (separate partitions)", code, r.runs())
}
}