From e6987f54b2b008f6a1eed2684c46911a0f6279c4 Mon Sep 17 00:00:00 2001 From: Steve Dudenhoeffer Date: Thu, 16 Jul 2026 16:54:30 -0400 Subject: [PATCH 1/2] =?UTF-8?q?feat:=20wave-3=20image=20+=20document=20sur?= =?UTF-8?q?faces=20=E2=80=94=20segmentation,=20colorize,=20face=20restore,?= =?UTF-8?q?=20ocr=20(ADR-0023)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - imagegen.Segmenter/SegmentationProvider: prompted mask via POST /upstream//v1/segment (file, prompt[, threshold], output=mask); white = prompted region, EditRequest.Mask polarity. - imagegen.Colorizer/ColorizeProvider: POST /upstream//v1/colorize. - imagegen.FaceRestorer/FaceRestoreProvider: POST /upstream//v1/restore_faces (upscale 1|2). - New ocr leaf package (Request/Page/Result, Recognize) + llamaswap OCRModel: POST /upstream//v1/ocr (file[, langs, max_pages]), tolerant per-page decode (join lines when page text absent), Raw escape hatch. - httptest contract tests per surface; ADR-0023; ADR index backfilled (0020-0022 rows were missing). Co-Authored-By: Claude Fable 5 --- docs/adr/0023-image-doc-surfaces.md | 59 +++++ docs/adr/README.md | 4 + imagegen/colorize.go | 54 +++++ imagegen/facerestore.go | 62 +++++ imagegen/segment.go | 70 ++++++ ocr/ocr.go | 114 +++++++++ provider/llamaswap/ocr.go | 140 +++++++++++ provider/llamaswap/ocr_test.go | 200 +++++++++++++++ provider/llamaswap/restore.go | 106 ++++++++ provider/llamaswap/segment.go | 76 ++++++ provider/llamaswap/segment_restore_test.go | 270 +++++++++++++++++++++ 11 files changed, 1155 insertions(+) create mode 100644 docs/adr/0023-image-doc-surfaces.md create mode 100644 imagegen/colorize.go create mode 100644 imagegen/facerestore.go create mode 100644 imagegen/segment.go create mode 100644 ocr/ocr.go create mode 100644 provider/llamaswap/ocr.go create mode 100644 provider/llamaswap/ocr_test.go create mode 100644 provider/llamaswap/restore.go create mode 100644 provider/llamaswap/segment.go create mode 100644 provider/llamaswap/segment_restore_test.go diff --git a/docs/adr/0023-image-doc-surfaces.md b/docs/adr/0023-image-doc-surfaces.md new file mode 100644 index 0000000..9bb48db --- /dev/null +++ b/docs/adr/0023-image-doc-surfaces.md @@ -0,0 +1,59 @@ +# ADR-0023: Wave-3 image + document surfaces (segmentation, colorize, face restore, OCR) + +Status: Accepted (2026-07-16) + +## Context + +The llama-swap host is gaining four wave-3 capabilities (spec: +mort docs/specs/2026-07-16-llamaswap-wave3.md): promptable segmentation +(GroundingDINO + SAM 2.1), photo colorization (DDColor), face restoration +(GFPGAN), and document OCR (Surya). All ride the `/upstream//` +passthrough (ADR-0020); none of their native APIs are OpenAI-shaped. + +## Decision + +1. **Grow `imagegen` with three optional interfaces**, ADR-0016→0022 + conventions (functional options, zero value = backend default, bytes-only + I/O, provider mints model): + - `imagegen.Segmenter` / `SegmentationProvider` — `SegmentationRequest + {Image, Prompt, Threshold}` → one grayscale mask image, WHITE = the + prompted region. Same polarity as `EditRequest.Mask` (white = repaint), + so the mask feeds inpainting directly; cutouts are derived client-side + (mask-as-alpha), one host call serving both. The client always sends + `output=mask` — the shim's `cutout`/`boxes` modes are not exposed. + - `imagegen.Colorizer` / `ColorizeProvider` — `ColorizeRequest{Image}`, + no knobs (the reference backend takes none; options reserved). + - `imagegen.FaceRestorer` / `FaceRestoreProvider` — + `FaceRestoreRequest{Image, Upscale}` (1 or 2, 0 = backend default). +2. **New `ocr` leaf package** rather than a method on an existing surface: + OCR is document-shaped (multi-page, PDFs), not image-generation-shaped. + `Request{Document, MIME, Filename, Languages, MaxPages}` → + `Result{Text, Pages []Page{Number, Text}, Raw}`. `Result.Text` is the + pages joined with a blank line; the per-line bbox/confidence/layout + detail stays in `Raw` (json.RawMessage) — exact text is the contract, + geometry is the escape hatch. +3. **provider/llamaswap wire shapes** (pinned by the netherstorm image + builds; host smoke tests are the drift defence): + - `POST /upstream//v1/segment` multipart `file`,`prompt` + [,`threshold`],`output=mask` → mask PNG. + - `POST /upstream//v1/colorize` multipart `file` → PNG. + - `POST /upstream//v1/restore_faces` multipart `file`[,`upscale`] → PNG. + - `POST /upstream//v1/ocr` multipart `file`[,`langs` (comma-joined), + `max_pages`] → JSON `{pages:[{number,text,lines,layout}]}`. The decode + is tolerant: a page without aggregate `text` joins its line texts; a + missing page `number` defaults to position. Zero pages is an error + (a blank page still arrives as a page), mirroring the "no transcript" + honesty rule. +4. **Binary success bodies are validated before wrapping** (ADR-0020 rule): + the three image surfaces reuse `singleImageResult` (positive evidence of + image-ness required), OCR requires decodable JSON. + +## Consequences + +- imagegen grows from five to eight optional surfaces; consumers + type-assert or use the provider methods directly, as before. +- `ocr` is the seventh leaf media package (imagegen, audio, videogen, + meshgen, musicgen, embeddings, ocr); the conventions have held across all + of them. +- PDF handling lives host-side (the shim rasterizes via pypdfium2); + majordomo ships bytes and never needs a PDF dependency. diff --git a/docs/adr/README.md b/docs/adr/README.md index 0c641e3..93772eb 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -23,3 +23,7 @@ One decision per file, append-only; supersede rather than rewrite. | [0017](0017-audio-interfaces.md) | audio — canonical speech synthesis + transcription interfaces | Accepted | | [0018](0018-imagegen-editor.md) | imagegen.Editor — image-to-image as a separate optional interface | Accepted | | [0019](0019-videogen-interface.md) | videogen — canonical video-generation surface | Accepted | +| [0020](0020-upstream-passthrough-media-surfaces.md) | Upstream-passthrough media surfaces (mask, upscale, background removal, interpolation, diarization, meshgen) | Accepted | +| [0021](0021-musicgen-interface.md) | musicgen — blocking Generate over an async job queue | Accepted | +| [0022](0022-embeddings-rerank-interface.md) | embeddings + rerank interface | Accepted | +| [0023](0023-image-doc-surfaces.md) | Wave-3 image + document surfaces (segmentation, colorize, face restore, OCR) | Accepted | diff --git a/imagegen/colorize.go b/imagegen/colorize.go new file mode 100644 index 0000000..0bad404 --- /dev/null +++ b/imagegen/colorize.go @@ -0,0 +1,54 @@ +package imagegen + +import "context" + +// ColorizeRequest asks a colorization backend (DDColor style) to add color to +// a grayscale/faded photo (ADR-0023). +type ColorizeRequest struct { + // Image is the image to colorize. Required. + Image Image +} + +// ColorizeOption mutates a ColorizeRequest before it is sent. Reserved for +// future request settings (the reference backend takes no parameters). +type ColorizeOption func(*ColorizeRequest) + +// Apply returns a copy of the request with all options applied. +func (r ColorizeRequest) Apply(opts ...ColorizeOption) ColorizeRequest { + for _, opt := range opts { + opt(&r) + } + return r +} + +// Colorizer adds color to grayscale images. +type Colorizer interface { + // Colorize returns the colorized image as a one-image Result. + Colorize(ctx context.Context, req ColorizeRequest, opts ...ColorizeOption) (*Result, error) +} + +// ColorizeModelOption configures a Colorizer at construction time. Reserved +// for future per-model settings. +type ColorizeModelOption func(*ColorizeModelConfig) + +// ColorizeModelConfig carries per-model construction settings. +type ColorizeModelConfig struct{} + +// ApplyColorizeModelOptions folds options into a config. +func ApplyColorizeModelOptions(opts []ColorizeModelOption) ColorizeModelConfig { + var cfg ColorizeModelConfig + for _, opt := range opts { + opt(&cfg) + } + return cfg +} + +// ColorizeProvider mints Colorizers bound to one backend. +type ColorizeProvider interface { + // Name is the registry identifier for the provider. + Name() string + + // ColorizeModel returns a Colorizer bound to the given id (passed through + // to the backend verbatim; no catalog validation). + ColorizeModel(id string, opts ...ColorizeModelOption) (Colorizer, error) +} diff --git a/imagegen/facerestore.go b/imagegen/facerestore.go new file mode 100644 index 0000000..7bc938a --- /dev/null +++ b/imagegen/facerestore.go @@ -0,0 +1,62 @@ +package imagegen + +import "context" + +// FaceRestoreRequest asks a face-restoration backend (GFPGAN style) to repair +// degraded faces in a photo. Zero values mean "backend default" (ADR-0023). +type FaceRestoreRequest struct { + // Image is the image to restore. Required. + Image Image + + // Upscale is the output enlargement factor (1 or 2 on the reference + // backend); 0 = backend default. + Upscale int +} + +// FaceRestoreOption mutates a FaceRestoreRequest before it is sent. +type FaceRestoreOption func(*FaceRestoreRequest) + +// WithFaceRestoreUpscale sets the output enlargement factor. +func WithFaceRestoreUpscale(n int) FaceRestoreOption { + return func(r *FaceRestoreRequest) { r.Upscale = n } +} + +// Apply returns a copy of the request with all options applied. +func (r FaceRestoreRequest) Apply(opts ...FaceRestoreOption) FaceRestoreRequest { + for _, opt := range opts { + opt(&r) + } + return r +} + +// FaceRestorer repairs degraded/blurry faces in photos. +type FaceRestorer interface { + // RestoreFaces returns the restored image as a one-image Result. + RestoreFaces(ctx context.Context, req FaceRestoreRequest, opts ...FaceRestoreOption) (*Result, error) +} + +// FaceRestoreModelOption configures a FaceRestorer at construction time. +// Reserved for future per-model settings. +type FaceRestoreModelOption func(*FaceRestoreModelConfig) + +// FaceRestoreModelConfig carries per-model construction settings. +type FaceRestoreModelConfig struct{} + +// ApplyFaceRestoreModelOptions folds options into a config. +func ApplyFaceRestoreModelOptions(opts []FaceRestoreModelOption) FaceRestoreModelConfig { + var cfg FaceRestoreModelConfig + for _, opt := range opts { + opt(&cfg) + } + return cfg +} + +// FaceRestoreProvider mints FaceRestorers bound to one backend. +type FaceRestoreProvider interface { + // Name is the registry identifier for the provider. + Name() string + + // FaceRestoreModel returns a FaceRestorer bound to the given id (passed + // through to the backend verbatim; no catalog validation). + FaceRestoreModel(id string, opts ...FaceRestoreModelOption) (FaceRestorer, error) +} diff --git a/imagegen/segment.go b/imagegen/segment.go new file mode 100644 index 0000000..e7ba12a --- /dev/null +++ b/imagegen/segment.go @@ -0,0 +1,70 @@ +package imagegen + +import "context" + +// SegmentationRequest asks a promptable-segmentation backend (GroundingDINO + +// SAM style) for the mask of a text-described region. Zero values mean +// "backend default" (ADR-0023). +type SegmentationRequest struct { + // Image is the image to segment. Required. + Image Image + + // Prompt is the text description of the region to segment (e.g. "the red + // car"). Required — promptless segmentation is a different capability. + Prompt string + + // Threshold is the detection confidence threshold in (0,1]; 0 = backend + // default. + Threshold float64 +} + +// SegmentationOption mutates a SegmentationRequest before it is sent. +type SegmentationOption func(*SegmentationRequest) + +// WithSegmentationThreshold sets the detection confidence threshold. +func WithSegmentationThreshold(t float64) SegmentationOption { + return func(r *SegmentationRequest) { r.Threshold = t } +} + +// Apply returns a copy of the request with all options applied. +func (r SegmentationRequest) Apply(opts ...SegmentationOption) SegmentationRequest { + for _, opt := range opts { + opt(&r) + } + return r +} + +// Segmenter extracts a prompted region's mask from an image. The result is a +// single grayscale mask image where WHITE marks the prompted region — the +// same polarity as EditRequest.Mask (white = repaint), so it feeds inpainting +// directly; derive a cutout client-side by applying the mask as alpha. +type Segmenter interface { + // Segment returns the region mask as a one-image Result. + Segment(ctx context.Context, req SegmentationRequest, opts ...SegmentationOption) (*Result, error) +} + +// SegmentationModelOption configures a Segmenter at construction time. +// Reserved for future per-model settings (mirrors ModelOption). +type SegmentationModelOption func(*SegmentationModelConfig) + +// SegmentationModelConfig carries per-model construction settings. +type SegmentationModelConfig struct{} + +// ApplySegmentationModelOptions folds options into a config. +func ApplySegmentationModelOptions(opts []SegmentationModelOption) SegmentationModelConfig { + var cfg SegmentationModelConfig + for _, opt := range opts { + opt(&cfg) + } + return cfg +} + +// SegmentationProvider mints Segmenters bound to one backend. +type SegmentationProvider interface { + // Name is the registry identifier for the provider. + Name() string + + // SegmentationModel returns a Segmenter bound to the given id (passed + // through to the backend verbatim; no catalog validation). + SegmentationModel(id string, opts ...SegmentationModelOption) (Segmenter, error) +} diff --git a/ocr/ocr.go b/ocr/ocr.go new file mode 100644 index 0000000..f07496b --- /dev/null +++ b/ocr/ocr.go @@ -0,0 +1,114 @@ +// Package ocr is majordomo's canonical document text-recognition surface. +// Like imagegen/audio/videogen/musicgen, it is a deliberately separate leaf +// contract from the llm package (ADR-0023, following the ADR-0016→0022 +// lineage: functional options, zero values = backend default, bytes-only +// I/O, Raw escape hatch). OCR is not chat-vision: it targets dedicated +// detection+recognition models (Surya style) that return exact per-page, +// per-line text rather than a model's paraphrase. +// +// The first implementation is provider/llamaswap, which posts the document +// to a Surya shim through the /upstream passthrough; the shim rasterizes +// PDFs itself, so Document may be an image (png/jpg/webp) or a PDF. +package ocr + +import "context" + +// Request is a text-recognition request. Zero values mean "backend default". +type Request struct { + // Document is the encoded document to recognize: an image (png/jpg/webp) + // or a PDF. Required. Carried as bytes, never a URL. + Document []byte + + // MIME is the document MIME type (e.g. "image/png", "application/pdf"); + // "" = let the backend sniff it. + MIME string + + // Filename is the multipart filename hint some backends key their format + // detection on; "" derives one from MIME ("document.pdf") or falls back + // to "document". + Filename string + + // Languages are ISO-639 hints for the recognizer (e.g. "en", "de"); + // nil = backend default (auto/multilingual). + Languages []string + + // MaxPages caps how many pages of a multi-page document are recognized; + // 0 = backend default (all pages, up to the backend's own ceiling). + MaxPages int +} + +// Page is the recognized text of one page. +type Page struct { + // Number is the 1-based page number. + Number int + + // Text is the page's recognized text. + Text string +} + +// Result is the canonical text-recognition result. +type Result struct { + // Text is the full recognized text, pages joined in order. + Text string + + // Pages are the per-page results in page order. + Pages []Page + + // Raw is the provider-native response object (e.g. the per-line + // bbox/confidence/layout detail), an escape hatch for provider-specific + // fields. May be nil; never required for normal use. + Raw any +} + +// Option mutates a Request before it is sent. Options passed to Recognize are +// applied to a copy of the request, so a Request value can be reused. +type Option func(*Request) + +// WithLanguages sets the recognizer language hints. +func WithLanguages(langs ...string) Option { + return func(r *Request) { r.Languages = langs } +} + +// WithMaxPages caps how many pages are recognized. +func WithMaxPages(n int) Option { return func(r *Request) { r.MaxPages = n } } + +// Apply returns a copy of the request with all options applied. Providers +// call this once at the top of Recognize. +func (r Request) Apply(opts ...Option) Request { + for _, opt := range opts { + opt(&r) + } + return r +} + +// Model recognizes text in documents. +type Model interface { + // Recognize extracts the document's text, page by page. + Recognize(ctx context.Context, req Request, opts ...Option) (*Result, error) +} + +// ModelOption configures a Model at construction time. Reserved for future +// per-model settings. +type ModelOption func(*ModelConfig) + +// ModelConfig carries per-model construction settings. +type ModelConfig struct{} + +// ApplyModelOptions folds options into a config. +func ApplyModelOptions(opts []ModelOption) ModelConfig { + var cfg ModelConfig + for _, opt := range opts { + opt(&cfg) + } + return cfg +} + +// Provider mints OCR models bound to one backend. +type Provider interface { + // Name is the registry identifier for the provider. + Name() string + + // OCRModel returns a Model bound to the given id (passed through to the + // backend verbatim; no catalog validation). + OCRModel(id string, opts ...ModelOption) (Model, error) +} diff --git a/provider/llamaswap/ocr.go b/provider/llamaswap/ocr.go new file mode 100644 index 0000000..2b7eb0b --- /dev/null +++ b/provider/llamaswap/ocr.go @@ -0,0 +1,140 @@ +// ocr.go implements ocr.Provider against a Surya-style shim reached through +// llama-swap's /upstream passthrough (ADR-0023): +// +// POST /upstream//v1/ocr multipart file[,langs,max_pages] +// +// The document may be an image or a PDF (the shim rasterizes PDFs itself). +// The response is per-page JSON: {pages:[{number,text,lines,layout}]}. The +// decode is tolerant — when a page carries no aggregate `text`, its line +// texts are joined instead — and the full payload survives in Result.Raw for +// callers that want bboxes/confidence/layout. +package llamaswap + +import ( + "context" + "encoding/json" + "fmt" + "mime" + "net/http" + "strconv" + "strings" + + "gitea.stevedudenhoeffer.com/steve/majordomo/llm" + "gitea.stevedudenhoeffer.com/steve/majordomo/ocr" +) + +// OCRModel implements ocr.Provider. The id selects which upstream llama-swap +// loads. +func (p *Provider) OCRModel(id string, opts ...ocr.ModelOption) (ocr.Model, error) { + if err := p.requireBaseURL(); err != nil { + return nil, err + } + _ = ocr.ApplyModelOptions(opts) + return &ocrModel{p: p, id: id}, nil +} + +type ocrModel struct { + p *Provider + id string +} + +// ocrResponse is the Surya shim's /v1/ocr shape (the subset this package +// relies on; per-line bbox/confidence and layout stay in Raw). +type ocrResponse struct { + Pages []struct { + Number int `json:"number"` + Text string `json:"text"` + Lines []struct { + Text string `json:"text"` + } `json:"lines"` + } `json:"pages"` +} + +// Recognize implements ocr.Model. +func (m *ocrModel) Recognize(ctx context.Context, req ocr.Request, opts ...ocr.Option) (*ocr.Result, error) { + req = req.Apply(opts...) + if len(req.Document) == 0 { + return nil, fmt.Errorf("%w: ocr requires document bytes", llm.ErrUnsupported) + } + if req.MaxPages < 0 { + return nil, fmt.Errorf("%w: ocr max pages must be >= 0, got %d", llm.ErrUnsupported, req.MaxPages) + } + path, err := upstreamPath(m.id, "/v1/ocr") + if err != nil { + return nil, err + } + maxPages := "" + if req.MaxPages != 0 { + maxPages = strconv.Itoa(req.MaxPages) + } + body, contentType, err := buildMultipart("build ocr form", + filePart{field: "file", filename: documentFilename(req.Filename, req.MIME), data: req.Document}, + []formField{ + {"langs", strings.Join(req.Languages, ","), false}, + {"max_pages", maxPages, false}, + }) + if err != nil { + return nil, err + } + raw, _, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxResponseBytes) + if err != nil { + return nil, err + } + + var out ocrResponse + if err := json.Unmarshal(raw, &out); err != nil { + return nil, fmt.Errorf("llama-swap: decode ocr response: %w", err) + } + if len(out.Pages) == 0 { + // A blank page still comes back as a page with empty text; zero pages + // is API drift or a soft error, never "the document was empty". + return nil, &llm.APIError{Provider: m.p.name, Model: m.id, + Message: "ocr response contained no pages: " + truncateForError(raw)} + } + res := &ocr.Result{Raw: json.RawMessage(raw)} + texts := make([]string, 0, len(out.Pages)) + for i, pg := range out.Pages { + text := pg.Text + if text == "" && len(pg.Lines) > 0 { + lines := make([]string, 0, len(pg.Lines)) + for _, ln := range pg.Lines { + lines = append(lines, ln.Text) + } + text = strings.Join(lines, "\n") + } + number := pg.Number + if number == 0 { + number = i + 1 + } + res.Pages = append(res.Pages, ocr.Page{Number: number, Text: text}) + texts = append(texts, text) + } + res.Text = strings.Join(texts, "\n\n") + return res, nil +} + +// documentFilename picks the multipart filename hint for an OCR document: the +// caller's (sanitized), else one derived from the MIME subtype +// ("document.pdf"), else "document". MIME parameters are stripped before +// matching, mirroring transcriptionFilename. +func documentFilename(filename, mimeType string) string { + if name := sanitizeFilename(filename); name != "" { + return name + } + mt := strings.ToLower(strings.TrimSpace(mimeType)) + if parsed, _, err := mime.ParseMediaType(mt); err == nil { + mt = parsed + } + switch mt { + case "application/pdf": + return "document.pdf" + case "image/png": + return "document.png" + case "image/jpeg", "image/jpg": + return "document.jpg" + case "image/webp": + return "document.webp" + default: + return "document" + } +} diff --git a/provider/llamaswap/ocr_test.go b/provider/llamaswap/ocr_test.go new file mode 100644 index 0000000..d6d1c38 --- /dev/null +++ b/provider/llamaswap/ocr_test.go @@ -0,0 +1,200 @@ +package llamaswap + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "gitea.stevedudenhoeffer.com/steve/majordomo/llm" + "gitea.stevedudenhoeffer.com/steve/majordomo/ocr" +) + +func TestRecognize(t *testing.T) { + var gotPath, gotLangs, gotMaxPages, gotFilename string + var gotFile []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + if err := r.ParseMultipartForm(1 << 20); err != nil { + t.Fatalf("parse multipart: %v", err) + } + gotLangs = r.FormValue("langs") + gotMaxPages = r.FormValue("max_pages") + f, hdr, err := r.FormFile("file") + if err != nil { + t.Fatalf("form file: %v", err) + } + defer f.Close() + gotFile, _ = io.ReadAll(f) + gotFilename = hdr.Filename + _, _ = w.Write([]byte(`{"pages":[ + {"number":1,"text":"page one","lines":[{"text":"page","bbox":[0,0,1,1],"confidence":0.9},{"text":"one"}],"layout":{}}, + {"number":2,"text":"page two"} + ]}`)) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + om, err := p.OCRModel("ocr-surya") + if err != nil { + t.Fatalf("OCRModel: %v", err) + } + res, err := om.Recognize(context.Background(), + ocr.Request{Document: []byte("%PDF"), MIME: "application/pdf"}, + ocr.WithLanguages("en", "de"), ocr.WithMaxPages(5)) + if err != nil { + t.Fatalf("Recognize: %v", err) + } + if gotPath != "/upstream/ocr-surya/v1/ocr" { + t.Errorf("path = %q", gotPath) + } + if gotLangs != "en,de" || gotMaxPages != "5" { + t.Errorf("langs/max_pages = %q/%q", gotLangs, gotMaxPages) + } + if string(gotFile) != "%PDF" || gotFilename != "document.pdf" { + t.Errorf("file = %q name = %q", gotFile, gotFilename) + } + if len(res.Pages) != 2 || res.Pages[0].Number != 1 || res.Pages[0].Text != "page one" || + res.Pages[1].Number != 2 || res.Pages[1].Text != "page two" { + t.Errorf("pages = %+v", res.Pages) + } + if res.Text != "page one\n\npage two" { + t.Errorf("text = %q", res.Text) + } + if res.Raw == nil { + t.Error("Raw = nil, want raw payload") + } +} + +func TestRecognizeJoinsLinesWhenPageTextAbsent(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{"pages":[{"lines":[{"text":"first line"},{"text":"second line"}]}]}`)) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + om, _ := p.OCRModel("ocr-surya") + res, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")}) + if err != nil { + t.Fatalf("Recognize: %v", err) + } + if len(res.Pages) != 1 || res.Pages[0].Text != "first line\nsecond line" { + t.Errorf("pages = %+v", res.Pages) + } + // Missing page number defaults to position. + if res.Pages[0].Number != 1 { + t.Errorf("number = %d, want 1", res.Pages[0].Number) + } + if res.Text != "first line\nsecond line" { + t.Errorf("text = %q", res.Text) + } +} + +func TestRecognizeOmitsUnsetFields(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseMultipartForm(1 << 20); err != nil { + t.Fatalf("parse multipart: %v", err) + } + for _, k := range []string{"langs", "max_pages"} { + if v, ok := r.MultipartForm.Value[k]; ok { + t.Errorf("unset request sent %q = %v, want omitted", k, v) + } + } + if _, hdr, err := r.FormFile("file"); err == nil { + if hdr.Filename != "document" { + t.Errorf("filename = %q, want document fallback", hdr.Filename) + } + } + _, _ = w.Write([]byte(`{"pages":[{"number":1,"text":"x"}]}`)) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + om, _ := p.OCRModel("ocr-surya") + if _, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")}); err != nil { + t.Fatalf("Recognize: %v", err) + } +} + +func TestRecognizeRejectsBadArgs(t *testing.T) { + p := New(WithBaseURL("http://unused")) + om, _ := p.OCRModel("ocr-surya") + if _, err := om.Recognize(context.Background(), ocr.Request{}); !errors.Is(err, llm.ErrUnsupported) { + t.Errorf("no document: err = %v, want ErrUnsupported", err) + } + if _, err := om.Recognize(context.Background(), + ocr.Request{Document: []byte{1}, MaxPages: -1}); !errors.Is(err, llm.ErrUnsupported) { + t.Errorf("negative max pages: err = %v, want ErrUnsupported", err) + } +} + +func TestRecognizeRejectsZeroPages(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{"pages":[]}`)) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + om, _ := p.OCRModel("ocr-surya") + _, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")}) + var apiErr *llm.APIError + if !errors.As(err, &apiErr) { + t.Fatalf("err = %v, want APIError for zero pages", err) + } +} + +func TestRecognizeRejectsNonJSONResponse(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + _, _ = w.Write([]byte("proxy error page")) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + om, _ := p.OCRModel("ocr-surya") + if _, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("img")}); err == nil || + !strings.Contains(err.Error(), "decode ocr response") { + t.Fatalf("err = %v, want decode error for non-JSON body", err) + } +} + +func TestRecognizeSurfacesAPIError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error":{"message":"unsupported file type"}}`)) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + om, _ := p.OCRModel("ocr-surya") + _, err := om.Recognize(context.Background(), ocr.Request{Document: []byte("bad")}) + var apiErr *llm.APIError + if !errors.As(err, &apiErr) { + t.Fatalf("err = %T %v, want *llm.APIError", err, err) + } + if apiErr.Status != http.StatusBadRequest || apiErr.Message != "unsupported file type" || apiErr.Model != "ocr-surya" { + t.Errorf("apiErr = %+v", apiErr) + } +} + +func TestDocumentFilename(t *testing.T) { + cases := []struct { + filename, mime, want string + }{ + {"scan.pdf", "", "scan.pdf"}, + {"evil\r\nname.pdf", "", "evilname.pdf"}, + {"", "application/pdf", "document.pdf"}, + {"", "image/png", "document.png"}, + {"", "image/jpeg", "document.jpg"}, + {"", "image/webp; charset=binary", "document.webp"}, + {"", "", "document"}, + } + for _, tc := range cases { + if got := documentFilename(tc.filename, tc.mime); got != tc.want { + t.Errorf("documentFilename(%q, %q) = %q, want %q", tc.filename, tc.mime, got, tc.want) + } + } +} diff --git a/provider/llamaswap/restore.go b/provider/llamaswap/restore.go new file mode 100644 index 0000000..e792e77 --- /dev/null +++ b/provider/llamaswap/restore.go @@ -0,0 +1,106 @@ +// restore.go implements the imagegen.ColorizeProvider and +// imagegen.FaceRestoreProvider surfaces against the mediautils shim reached +// through llama-swap's /upstream passthrough (ADR-0023): +// +// colorize POST /upstream//v1/colorize (DDColor) +// restore_faces POST /upstream//v1/restore_faces (GFPGAN) +// +// Both are one-file multipart in, one PNG out, mirroring mediautil.go. +package llamaswap + +import ( + "context" + "fmt" + "net/http" + "strconv" + + "gitea.stevedudenhoeffer.com/steve/majordomo/imagegen" + "gitea.stevedudenhoeffer.com/steve/majordomo/llm" +) + +// --- colorize --- + +// ColorizeModel implements imagegen.ColorizeProvider against the mediautils +// shim's POST /v1/colorize. The id selects which upstream llama-swap loads. +func (p *Provider) ColorizeModel(id string, opts ...imagegen.ColorizeModelOption) (imagegen.Colorizer, error) { + if err := p.requireBaseURL(); err != nil { + return nil, err + } + _ = imagegen.ApplyColorizeModelOptions(opts) + return &colorizeModel{p: p, id: id}, nil +} + +type colorizeModel struct { + p *Provider + id string +} + +// Colorize implements imagegen.Colorizer. +func (m *colorizeModel) Colorize(ctx context.Context, req imagegen.ColorizeRequest, opts ...imagegen.ColorizeOption) (*imagegen.Result, error) { + req = req.Apply(opts...) + if len(req.Image.Data) == 0 { + return nil, fmt.Errorf("%w: colorization requires an image", llm.ErrUnsupported) + } + path, err := upstreamPath(m.id, "/v1/colorize") + if err != nil { + return nil, err + } + body, contentType, err := buildMultipart("build colorize form", + filePart{field: "file", filename: "image.png", data: req.Image.Data}, + nil) + if err != nil { + return nil, err + } + raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxImageResponseBytes) + if err != nil { + return nil, err + } + return singleImageResult(m.p.name, m.id, "colorize", raw, respType) +} + +// --- face restoration --- + +// FaceRestoreModel implements imagegen.FaceRestoreProvider against the +// mediautils shim's POST /v1/restore_faces. +func (p *Provider) FaceRestoreModel(id string, opts ...imagegen.FaceRestoreModelOption) (imagegen.FaceRestorer, error) { + if err := p.requireBaseURL(); err != nil { + return nil, err + } + _ = imagegen.ApplyFaceRestoreModelOptions(opts) + return &faceRestoreModel{p: p, id: id}, nil +} + +type faceRestoreModel struct { + p *Provider + id string +} + +// RestoreFaces implements imagegen.FaceRestorer. +func (m *faceRestoreModel) RestoreFaces(ctx context.Context, req imagegen.FaceRestoreRequest, opts ...imagegen.FaceRestoreOption) (*imagegen.Result, error) { + req = req.Apply(opts...) + if len(req.Image.Data) == 0 { + return nil, fmt.Errorf("%w: face restoration requires an image", llm.ErrUnsupported) + } + if req.Upscale != 0 && req.Upscale != 1 && req.Upscale != 2 { + return nil, fmt.Errorf("%w: face-restore upscale must be 1 or 2, got %d", llm.ErrUnsupported, req.Upscale) + } + path, err := upstreamPath(m.id, "/v1/restore_faces") + if err != nil { + return nil, err + } + upscale := "" + if req.Upscale != 0 { + upscale = strconv.Itoa(req.Upscale) + } + body, contentType, err := buildMultipart("build restore-faces form", + filePart{field: "file", filename: "image.png", data: req.Image.Data}, + []formField{{"upscale", upscale, false}}) + if err != nil { + return nil, err + } + raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxImageResponseBytes) + if err != nil { + return nil, err + } + return singleImageResult(m.p.name, m.id, "face restoration", raw, respType) +} diff --git a/provider/llamaswap/segment.go b/provider/llamaswap/segment.go new file mode 100644 index 0000000..404d592 --- /dev/null +++ b/provider/llamaswap/segment.go @@ -0,0 +1,76 @@ +// segment.go implements imagegen.SegmentationProvider against a +// GroundingDINO+SAM shim (segment-langsam) reached through llama-swap's +// /upstream passthrough (ADR-0023): +// +// POST /upstream//v1/segment multipart file,prompt[,threshold],output=mask +// +// The response is a single grayscale mask PNG where WHITE marks the prompted +// region — directly usable as imagegen.EditRequest.Mask (white = repaint). +// The shim also offers output=cutout|boxes; this client always requests the +// mask, because a cutout is derivable client-side from mask+original with no +// second GPU call. +package llamaswap + +import ( + "context" + "fmt" + "net/http" + "strconv" + + "gitea.stevedudenhoeffer.com/steve/majordomo/imagegen" + "gitea.stevedudenhoeffer.com/steve/majordomo/llm" +) + +// SegmentationModel implements imagegen.SegmentationProvider. The id selects +// which upstream llama-swap loads. +func (p *Provider) SegmentationModel(id string, opts ...imagegen.SegmentationModelOption) (imagegen.Segmenter, error) { + if err := p.requireBaseURL(); err != nil { + return nil, err + } + _ = imagegen.ApplySegmentationModelOptions(opts) + return &segmentationModel{p: p, id: id}, nil +} + +type segmentationModel struct { + p *Provider + id string +} + +// Segment implements imagegen.Segmenter. +func (m *segmentationModel) Segment(ctx context.Context, req imagegen.SegmentationRequest, opts ...imagegen.SegmentationOption) (*imagegen.Result, error) { + req = req.Apply(opts...) + if len(req.Image.Data) == 0 { + return nil, fmt.Errorf("%w: segmentation requires an image", llm.ErrUnsupported) + } + if req.Prompt == "" { + return nil, fmt.Errorf("%w: segmentation requires a prompt", llm.ErrUnsupported) + } + if req.Threshold < 0 || req.Threshold > 1 { + return nil, fmt.Errorf("%w: segmentation threshold must be in [0,1], got %g", llm.ErrUnsupported, req.Threshold) + } + path, err := upstreamPath(m.id, "/v1/segment") + if err != nil { + return nil, err + } + threshold := "" + if req.Threshold != 0 { + threshold = strconv.FormatFloat(req.Threshold, 'g', -1, 64) + } + body, contentType, err := buildMultipart("build segment form", + filePart{field: "file", filename: "image.png", data: req.Image.Data}, + []formField{ + {"prompt", req.Prompt, true}, + {"threshold", threshold, false}, + // Always the mask: white = prompted region, EditRequest.Mask + // polarity. Cutouts are derived client-side. + {"output", "mask", true}, + }) + if err != nil { + return nil, err + } + raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxImageResponseBytes) + if err != nil { + return nil, err + } + return singleImageResult(m.p.name, m.id, "segmentation", raw, respType) +} diff --git a/provider/llamaswap/segment_restore_test.go b/provider/llamaswap/segment_restore_test.go new file mode 100644 index 0000000..fa39d91 --- /dev/null +++ b/provider/llamaswap/segment_restore_test.go @@ -0,0 +1,270 @@ +package llamaswap + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "testing" + + "gitea.stevedudenhoeffer.com/steve/majordomo/imagegen" + "gitea.stevedudenhoeffer.com/steve/majordomo/llm" +) + +func TestSegment(t *testing.T) { + png := pngFixture(t) + var gotPath, gotPrompt, gotThreshold, gotOutput, gotFilename string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + if err := r.ParseMultipartForm(1 << 20); err != nil { + t.Errorf("parse multipart: %v", err) + } + gotPrompt = r.FormValue("prompt") + gotThreshold = r.FormValue("threshold") + gotOutput = r.FormValue("output") + if f, hdr, err := r.FormFile("file"); err == nil { + gotFilename = hdr.Filename + f.Close() + } else { + t.Errorf("file part: %v", err) + } + w.Header().Set("Content-Type", "image/png") + _, _ = w.Write(png) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + sg, err := p.SegmentationModel("segment-langsam") + if err != nil { + t.Fatalf("SegmentationModel: %v", err) + } + res, err := sg.Segment(context.Background(), + imagegen.SegmentationRequest{Image: imagegen.Image{MIME: "image/png", Data: png}, Prompt: "the red car"}, + imagegen.WithSegmentationThreshold(0.35)) + if err != nil { + t.Fatalf("Segment: %v", err) + } + if gotPath != "/upstream/segment-langsam/v1/segment" { + t.Errorf("path = %q", gotPath) + } + if gotPrompt != "the red car" || gotThreshold != "0.35" || gotOutput != "mask" || gotFilename != "image.png" { + t.Errorf("prompt/threshold/output/filename = %q/%q/%q/%q", gotPrompt, gotThreshold, gotOutput, gotFilename) + } + if len(res.Images) != 1 || res.Images[0].MIME != "image/png" { + t.Fatalf("images = %+v", res.Images) + } +} + +func TestSegmentOmitsDefaultThreshold(t *testing.T) { + png := pngFixture(t) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseMultipartForm(1 << 20); err != nil { + t.Errorf("parse multipart: %v", err) + } + if _, ok := r.MultipartForm.Value["threshold"]; ok { + t.Error("threshold field sent for default; want omitted") + } + if got := r.FormValue("output"); got != "mask" { + t.Errorf("output = %q, want mask (always sent)", got) + } + w.Header().Set("Content-Type", "image/png") + _, _ = w.Write(png) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + sg, _ := p.SegmentationModel("segment-langsam") + if _, err := sg.Segment(context.Background(), + imagegen.SegmentationRequest{Image: imagegen.Image{Data: png}, Prompt: "dog"}); err != nil { + t.Fatalf("Segment: %v", err) + } +} + +func TestSegmentRejectsBadArgs(t *testing.T) { + p := New(WithBaseURL("http://unused")) + sg, _ := p.SegmentationModel("segment-langsam") + if _, err := sg.Segment(context.Background(), + imagegen.SegmentationRequest{Prompt: "dog"}); !errors.Is(err, llm.ErrUnsupported) { + t.Errorf("no image: err = %v, want ErrUnsupported", err) + } + if _, err := sg.Segment(context.Background(), + imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}}); !errors.Is(err, llm.ErrUnsupported) { + t.Errorf("no prompt: err = %v, want ErrUnsupported", err) + } + if _, err := sg.Segment(context.Background(), + imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}, Prompt: "dog", Threshold: 1.5}); !errors.Is(err, llm.ErrUnsupported) { + t.Errorf("threshold 1.5: err = %v, want ErrUnsupported", err) + } +} + +func TestSegmentRejectsNonImageResponse(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + _, _ = w.Write([]byte("proxy error page")) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + sg, _ := p.SegmentationModel("segment-langsam") + _, err := sg.Segment(context.Background(), + imagegen.SegmentationRequest{Image: imagegen.Image{Data: pngFixture(t)}, Prompt: "dog"}) + var apiErr *llm.APIError + if !errors.As(err, &apiErr) { + t.Fatalf("err = %v, want APIError for non-image body", err) + } +} + +func TestColorize(t *testing.T) { + png := pngFixture(t) + var gotPath, gotFilename string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + if err := r.ParseMultipartForm(1 << 20); err != nil { + t.Errorf("parse multipart: %v", err) + } + if f, hdr, err := r.FormFile("file"); err == nil { + gotFilename = hdr.Filename + f.Close() + } else { + t.Errorf("file part: %v", err) + } + w.Header().Set("Content-Type", "image/png") + _, _ = w.Write(png) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + cl, err := p.ColorizeModel("mediautils") + if err != nil { + t.Fatalf("ColorizeModel: %v", err) + } + res, err := cl.Colorize(context.Background(), + imagegen.ColorizeRequest{Image: imagegen.Image{Data: png}}) + if err != nil { + t.Fatalf("Colorize: %v", err) + } + if gotPath != "/upstream/mediautils/v1/colorize" { + t.Errorf("path = %q", gotPath) + } + if gotFilename != "image.png" { + t.Errorf("filename = %q", gotFilename) + } + if len(res.Images) != 1 || res.Images[0].MIME != "image/png" { + t.Fatalf("images = %+v", res.Images) + } +} + +func TestColorizeRejectsEmptyImage(t *testing.T) { + p := New(WithBaseURL("http://unused")) + cl, _ := p.ColorizeModel("mediautils") + if _, err := cl.Colorize(context.Background(), imagegen.ColorizeRequest{}); !errors.Is(err, llm.ErrUnsupported) { + t.Errorf("err = %v, want ErrUnsupported", err) + } +} + +func TestColorizeSurfacesAPIError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusServiceUnavailable) + _, _ = w.Write([]byte(`{"error":{"message":"model is loading"}}`)) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + cl, _ := p.ColorizeModel("mediautils") + _, err := cl.Colorize(context.Background(), + imagegen.ColorizeRequest{Image: imagegen.Image{Data: pngFixture(t)}}) + var apiErr *llm.APIError + if !errors.As(err, &apiErr) { + t.Fatalf("err = %T %v, want *llm.APIError", err, err) + } + if apiErr.Status != http.StatusServiceUnavailable || apiErr.Message != "model is loading" || apiErr.Model != "mediautils" { + t.Errorf("apiErr = %+v", apiErr) + } +} + +func TestRestoreFaces(t *testing.T) { + png := pngFixture(t) + var gotPath, gotUpscale string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + if err := r.ParseMultipartForm(1 << 20); err != nil { + t.Errorf("parse multipart: %v", err) + } + gotUpscale = r.FormValue("upscale") + w.Header().Set("Content-Type", "image/png") + _, _ = w.Write(png) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + fr, err := p.FaceRestoreModel("mediautils") + if err != nil { + t.Fatalf("FaceRestoreModel: %v", err) + } + res, err := fr.RestoreFaces(context.Background(), + imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: png}}, + imagegen.WithFaceRestoreUpscale(2)) + if err != nil { + t.Fatalf("RestoreFaces: %v", err) + } + if gotPath != "/upstream/mediautils/v1/restore_faces" { + t.Errorf("path = %q", gotPath) + } + if gotUpscale != "2" { + t.Errorf("upscale = %q", gotUpscale) + } + if len(res.Images) != 1 { + t.Fatalf("images = %+v", res.Images) + } +} + +func TestRestoreFacesOmitsDefaultUpscale(t *testing.T) { + png := pngFixture(t) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseMultipartForm(1 << 20); err != nil { + t.Errorf("parse multipart: %v", err) + } + if _, ok := r.MultipartForm.Value["upscale"]; ok { + t.Error("upscale field sent for default; want omitted") + } + w.Header().Set("Content-Type", "image/png") + _, _ = w.Write(png) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + fr, _ := p.FaceRestoreModel("mediautils") + if _, err := fr.RestoreFaces(context.Background(), + imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: png}}); err != nil { + t.Fatalf("RestoreFaces: %v", err) + } +} + +func TestRestoreFacesRejectsBadArgs(t *testing.T) { + p := New(WithBaseURL("http://unused")) + fr, _ := p.FaceRestoreModel("mediautils") + if _, err := fr.RestoreFaces(context.Background(), imagegen.FaceRestoreRequest{}); !errors.Is(err, llm.ErrUnsupported) { + t.Errorf("no image: err = %v, want ErrUnsupported", err) + } + if _, err := fr.RestoreFaces(context.Background(), + imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: []byte{1}}, Upscale: 3}); !errors.Is(err, llm.ErrUnsupported) { + t.Errorf("upscale 3: err = %v, want ErrUnsupported", err) + } +} + +func TestRestoreFacesRejectsNonImageResponse(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header()["Content-Type"] = nil // NO Content-Type at all + _, _ = w.Write([]byte("502 bad gateway")) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + fr, _ := p.FaceRestoreModel("mediautils") + _, err := fr.RestoreFaces(context.Background(), + imagegen.FaceRestoreRequest{Image: imagegen.Image{Data: pngFixture(t)}}) + var apiErr *llm.APIError + if !errors.As(err, &apiErr) { + t.Fatalf("err = %v, want APIError for headerless non-image body", err) + } +} -- 2.54.0 From 966ea16166696f342a6cce5309aeaf99f18bd681 Mon Sep 17 00:00:00 2001 From: Steve Dudenhoeffer Date: Thu, 16 Jul 2026 19:00:11 -0400 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20review=20findings=20=E2=80=94=20NaN?= =?UTF-8?q?=20threshold=20guard,=20percent-escape=20rejection=20in=20upstr?= =?UTF-8?q?eam=20model=20ids?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Segment: reject NaN thresholds (NaN fails every comparison, so it passed the [0,1] range check and reached the shim as the literal string "NaN"); ±Inf were already caught by the range comparisons, now covered by tests too. - upstreamPath: reject '%' in model ids — %2F/%2E%2E percent-escapes decode back into path structure server-side, bypassing the literal /?#/.. rejection. Ids never legitimately contain '%'. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WWCQcYStWXBYUy5sZWnbLT --- provider/llamaswap/mediautil_test.go | 2 +- provider/llamaswap/segment.go | 5 ++++- provider/llamaswap/segment_restore_test.go | 11 +++++++++++ provider/llamaswap/upstream.go | 6 ++++-- 4 files changed, 20 insertions(+), 4 deletions(-) diff --git a/provider/llamaswap/mediautil_test.go b/provider/llamaswap/mediautil_test.go index 466ec9e..9bb42c3 100644 --- a/provider/llamaswap/mediautil_test.go +++ b/provider/llamaswap/mediautil_test.go @@ -238,7 +238,7 @@ func TestInterpolateRejectsBadArgs(t *testing.T) { } func TestUpstreamPathRejectsSeparators(t *testing.T) { - for _, bad := range []string{"", "a/b", "a?b", "a#b"} { + for _, bad := range []string{"", "a/b", "a?b", "a#b", "a%2Fb", "%2e%2e", "a%b"} { if _, err := upstreamPath(bad, "/x"); err == nil { t.Errorf("upstreamPath(%q) succeeded; want error", bad) } diff --git a/provider/llamaswap/segment.go b/provider/llamaswap/segment.go index 404d592..c92123f 100644 --- a/provider/llamaswap/segment.go +++ b/provider/llamaswap/segment.go @@ -14,6 +14,7 @@ package llamaswap import ( "context" "fmt" + "math" "net/http" "strconv" @@ -45,7 +46,9 @@ func (m *segmentationModel) Segment(ctx context.Context, req imagegen.Segmentati if req.Prompt == "" { return nil, fmt.Errorf("%w: segmentation requires a prompt", llm.ErrUnsupported) } - if req.Threshold < 0 || req.Threshold > 1 { + // NaN fails every comparison, so it would sail through a bare range + // check and reach the upstream as the literal string "NaN". + if math.IsNaN(req.Threshold) || req.Threshold < 0 || req.Threshold > 1 { return nil, fmt.Errorf("%w: segmentation threshold must be in [0,1], got %g", llm.ErrUnsupported, req.Threshold) } path, err := upstreamPath(m.id, "/v1/segment") diff --git a/provider/llamaswap/segment_restore_test.go b/provider/llamaswap/segment_restore_test.go index fa39d91..dcb4f21 100644 --- a/provider/llamaswap/segment_restore_test.go +++ b/provider/llamaswap/segment_restore_test.go @@ -3,6 +3,7 @@ package llamaswap import ( "context" "errors" + "math" "net/http" "net/http/httptest" "testing" @@ -95,6 +96,16 @@ func TestSegmentRejectsBadArgs(t *testing.T) { imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}, Prompt: "dog", Threshold: 1.5}); !errors.Is(err, llm.ErrUnsupported) { t.Errorf("threshold 1.5: err = %v, want ErrUnsupported", err) } + for name, bad := range map[string]float64{ + "NaN": math.NaN(), + "+Inf": math.Inf(1), + "-Inf": math.Inf(-1), + } { + if _, err := sg.Segment(context.Background(), + imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}, Prompt: "dog", Threshold: bad}); !errors.Is(err, llm.ErrUnsupported) { + t.Errorf("threshold %s: err = %v, want ErrUnsupported", name, err) + } + } } func TestSegmentRejectsNonImageResponse(t *testing.T) { diff --git a/provider/llamaswap/upstream.go b/provider/llamaswap/upstream.go index 1a33bc2..48998d2 100644 --- a/provider/llamaswap/upstream.go +++ b/provider/llamaswap/upstream.go @@ -16,12 +16,14 @@ import ( // // Why reject rather than escape: same rationale as Unload — model ids // legitimately contain ":" but never path-structure characters, and escaping -// would mask a config error instead of surfacing it. +// would mask a config error instead of surfacing it. '%' is rejected too: +// ids never legitimately carry percent-escapes, and %2F/%2E%2E would decode +// back into path structure on the server side. func upstreamPath(model, rest string) (string, error) { if strings.TrimSpace(model) == "" { return "", fmt.Errorf("llama-swap: upstream call requires a model id") } - if strings.ContainsAny(model, "/?#") || strings.Contains(model, "..") { + if strings.ContainsAny(model, "/?#%") || strings.Contains(model, "..") { return "", fmt.Errorf("llama-swap: invalid model id %q for upstream call (contains a path separator)", model) } if !strings.HasPrefix(rest, "/") { -- 2.54.0