LerNotaFiscal/app/ai_extraction.py

"""Extração de documentos fiscais via OpenAI Vision, com fallback local.

`extract_with_ai` recebe as imagens de um arquivo enviado (páginas de PDF já
renderizadas em PNG, ou a própria imagem) e devolve uma lista de `RawExtraction`
— ou `None` para sinalizar que o chamador deve cair no OCR/heurística local.

A superfície da SDK usada (chat.completions + response_format json_object +
entrada de imagem por data URI) foi verificada contra `openai` 2.x.
"""

from __future__ import annotations

import base64
import json
import logging
from dataclasses import dataclass, field
from pathlib import Path

from .config import get_settings

logger = logging.getLogger("lernotafiscal.ai")


@dataclass
class RawExtraction:
    """Resultado bruto de UM documento fiscal, antes de normalizar a data."""

    supplier_name: str | None
    purchase_date_raw: str | None
    total_paid: float | None
    legible: bool
    uncertain_fields: list[str] = field(default_factory=list)
    extractor: str = "openai"
    raw_text: str = ""


_MIME = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg"}

_SYSTEM_PROMPT = (
    "Você é um extrator de dados de notas e cupons fiscais brasileiros. "
    "Responda SEMPRE em JSON válido, sem texto fora do JSON."
)

_USER_PROMPT = """As imagens a seguir são documentos fiscais (notas, cupons, recibos). Podem conter mais de um documento.
Para CADA documento distinto, extraia:
- "fornecedor": razão social ou nome fantasia do emissor (o mais destacado no cabeçalho), ou null.
- "data_compra": data da compra no formato "YYYY-MM-DD", ou null se ilegível/ausente.
- "valor_pago": número (ponto decimal) do TOTAL efetivamente pago. Use o "VALOR TOTAL"/"TOTAL A PAGAR"/"VALOR PAGO". NUNCA use "TROCO" nem "DINHEIRO RECEBIDO". Se parcelado, use o total da compra. Null se ilegível.
- "legivel": true se você leu os campos com confiança; false se o documento está borrado, cortado ou ilegível.
- "campos_incertos": lista dos campos que ficaram duvidosos (ex.: ["data_compra","valor_pago"]).

Responda exatamente neste formato:
{"documentos": [{"fornecedor": ..., "data_compra": ..., "valor_pago": ..., "legivel": ..., "campos_incertos": [...]}]}
Se nenhum documento fiscal for identificável, responda {"documentos": []}."""


def _data_uri(path: Path) -> str | None:
    mime = _MIME.get(path.suffix.lower())
    if not mime:
        return None
    try:
        encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    except OSError:
        return None
    return f"data:{mime};base64,{encoded}"


def _coerce_float(value: object) -> float | None:
    if value is None:
        return None
    if isinstance(value, (int, float)):
        return round(float(value), 2)
    text = str(value).strip().replace("R$", "").replace(" ", "")
    if not text:
        return None
    # aceita "1.234,56" e "1234.56"
    if "," in text and "." in text:
        text = text.replace(".", "").replace(",", ".")
    elif "," in text:
        text = text.replace(",", ".")
    try:
        return round(float(text), 2)
    except ValueError:
        return None


def _parse_response(content: str) -> list[RawExtraction]:
    data = json.loads(content)
    docs = data.get("documentos") if isinstance(data, dict) else None
    if not isinstance(docs, list):
        return []
    results: list[RawExtraction] = []
    for item in docs:
        if not isinstance(item, dict):
            continue
        supplier = item.get("fornecedor")
        supplier = str(supplier).strip()[:200] if supplier else None
        date_raw = item.get("data_compra")
        date_raw = str(date_raw).strip() if date_raw else None
        total = _coerce_float(item.get("valor_pago"))
        legible = bool(item.get("legivel", True))
        uncertain = item.get("campos_incertos") or []
        uncertain = [str(x) for x in uncertain if x] if isinstance(uncertain, list) else []
        # segurança: campos faltando são inerentemente incertos
        for name, val in (("fornecedor", supplier), ("data_compra", date_raw), ("valor_pago", total)):
            if val is None and name not in uncertain:
                uncertain.append(name)
        if uncertain and legible and len(uncertain) >= 2:
            legible = False
        results.append(
            RawExtraction(
                supplier_name=supplier,
                purchase_date_raw=date_raw,
                total_paid=total,
                legible=legible,
                uncertain_fields=uncertain,
            )
        )
    return results


def extract_with_ai(image_paths: list[Path]) -> list[RawExtraction] | None:
    """Extrai via OpenAI. Retorna None se IA indisponível ou em erro (=> fallback)."""
    settings = get_settings()
    if not settings.openai_enabled:
        return None

    uris = [uri for p in image_paths[: settings.openai_max_pages] if (uri := _data_uri(p))]
    if not uris:
        return None

    try:
        from openai import OpenAI

        client = OpenAI(api_key=settings.openai_api_key)
        content: list[dict] = [{"type": "text", "text": _USER_PROMPT}]
        for uri in uris:
            content.append({"type": "image_url", "image_url": {"url": uri}})

        response = client.chat.completions.create(
            model=settings.openai_model,
            messages=[
                {"role": "system", "content": _SYSTEM_PROMPT},
                {"role": "user", "content": content},
            ],
            response_format={"type": "json_object"},
            temperature=0,
        )
        message = response.choices[0].message.content or "{}"
        return _parse_response(message)
    except Exception as exc:  # rede, cota, parsing, modelo indisponível...
        logger.warning("Extração OpenAI falhou, usando fallback local: %s", exc)
        return None