File size: 2,307 Bytes
2edb151 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | 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)
|