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,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