LerNota/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 datetime import date
from pathlib import Path

from .config import get_openai_client, 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."
)

# Tolerância de plausibilidade da data extraída (ver `_is_plausible_purchase_date`).
# A checagem opera em granularidade de mês/ano — é isso que de fato importa
# para a competência persistida, o dia é descartado na normalização.
_DATE_TOLERANCE_YEARS_PAST = 2
_DATE_TOLERANCE_MONTHS_FUTURE = 1


def _user_prompt() -> str:
    today = date.today().isoformat()
    return f"""A data de hoje é {today}. 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. Esses documentos são quase sempre recentes (deste ano ou do ano anterior a {today[:4]}) — releia com cuidado os dois últimos dígitos do ano antes de responder, para não confundir dígitos parecidos (ex.: não troque "26" por "22").
- "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 _is_plausible_purchase_date(date_raw: str) -> bool:
    """Sanidade sobre a competência (mês/ano) da data devolvida pela IA: notas
    fiscais são quase sempre recentes, então uma competência muito no passado
    ou no futuro é sinal de erro de leitura de dígito (ex.: "26" lido como
    "22") — melhor mandar para revisão manual do que aceitar silenciosamente."""
    try:
        year, month, day = (int(part) for part in date_raw.split("-"))
        date(year, month, day)  # valida o calendário
    except (ValueError, TypeError):
        return True  # formato inesperado já é tratado como campo incerto à parte
    today = date.today()
    months_ahead = (year * 12 + month) - (today.year * 12 + today.month)
    if months_ahead > _DATE_TOLERANCE_MONTHS_FUTURE:
        return False
    if year < today.year - _DATE_TOLERANCE_YEARS_PAST:
        return False
    return True


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)
        # segurança: data implausível (ano muito no passado/futuro) força revisão
        # manual, mesmo que a IA tenha respondido "legivel": true — ver
        # `_is_plausible_purchase_date` para o porquê (erro de leitura de dígito).
        if date_raw is not None and not _is_plausible_purchase_date(date_raw):
            if "data_compra" not in uncertain:
                uncertain.append("data_compra")
            legible = False
        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:
        client = get_openai_client()
        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