api: serve PDFs — extraction in the child under the pool, full text or map, errors
Route "pdf" through clean.Kind in doFetch/cleanFetched: extract under the pool (keyed by caller class + document hash, one per caller), map the five pdftext errors plus a pool panic to 422, ErrBusy to 503+Retry-After 10, ErrCallerBusy to 429, and answer with the whole text or the map. Adds RetryAfter to fetch.Error and the matching header. No cache, pages or grep — those are task 07. Co-Authored-By: ornith-1.5-35b-a3b <ornith-1.5-35b-a3b@llama.cpp.invalid>
This commit is contained in:
ornith-1.5-35b-a3b
@@ -0,0 +1,402 @@
|
||||
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)
|
||||
}
|
||||
if DefaultPDFMaxQuestionChunks != 200 {
|
||||
t.Errorf("DefaultPDFMaxQuestionChunks = %d, want 200", DefaultPDFMaxQuestionChunks)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user