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:
@@ -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.
|
||||||
@@ -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.
|
||||||
@@ -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.
|
||||||
@@ -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))
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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.
|
||||||
@@ -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.
|
||||||
@@ -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())
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user