app/extract.py

from __future__ import annotations

from app.config import CATEGORIES, DOC_KINDS, Settings
from app.schemas import ReceiptExtract, parse_extract_json
from backends.base import LLMBackend

SYSTEM = (
    "You extract structured data from a photo of a receipt, invoice, or paper document. "
    "Reply with a single JSON object only — no markdown, no commentary, no trailing text."
)

_SCHEMA_HINT = f"""
Required keys:
  doc_kind: one of {list(DOC_KINDS)}
  category: one of {list(CATEGORIES)}
  vendor: string or null
  date: YYYY-MM-DD or null
  tax: number or null
  total: number or null
  currency: string or null (ISO 4217 if known)
  line_items: array of objects with description (string), qty (number|null),
              unit_price (number|null), amount (number|null), sku (string|null)

Rules:
- Money as numbers, not strings. Unknown fields must be null.
- Do not invent SKUs, vendors, or totals. If unreadable, use null.
- Prefer the printed total over summing line items when they disagree.
- category is the spend bucket (groceries, dining, …), not the store name.
""".strip()


def build_user_prompt(*, ocr_text: str | None, hint: str | None = None) -> str:
    parts = [_SCHEMA_HINT]
    if ocr_text:
        parts.append("OCR assist (may be noisy):\n" + ocr_text.strip()[:8000])
    if hint:
        parts.append(hint)
    parts.append("Extract the JSON now.")
    return "\n\n".join(parts)


def extract_receipt(
    llm: LLMBackend,
    *,
    settings: Settings,
    image_jpeg: bytes | None,
    ocr_text: str | None,
) -> ReceiptExtract:
    del settings
    if image_jpeg and not llm.accepts_images:
        image_jpeg = None
    if image_jpeg is None and not (ocr_text and ocr_text.strip()):
        raise ValueError("need an image (vision LLM) or OCR/text to extract")
    user = build_user_prompt(ocr_text=ocr_text)
    raw = llm.complete_json(system=SYSTEM, user=user, image_jpeg=image_jpeg)
    try:
        return parse_extract_json(raw)
    except (ValueError, Exception) as first:
        retry = build_user_prompt(
            ocr_text=ocr_text,
            hint=f"Previous output failed validation: {first}. Return corrected JSON only.",
        )
        raw2 = llm.complete_json(system=SYSTEM, user=retry, image_jpeg=image_jpeg)
        return parse_extract_json(raw2)