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>
2026-09-26 11:01:18 -07:00
parent 1311cf38d4
commit 49d96e197b
14 changed files with 1464 additions and 0 deletions
@@ -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())
+	}
+}