majordomo/provider/llamaswap/ocr.go

// ocr.go implements ocr.Provider against a Surya-style shim reached through
// llama-swap's /upstream passthrough (ADR-0023):
//
//	POST /upstream/<id>/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"
	}
}