nia-etl-utils 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

nia_etl_utils/ocr.py

@@ -0,0 +1,327 @@
+"""Módulo para processamento de OCR via API IntelliDoc.
+
+Este módulo fornece funções para submeter documentos ao serviço de OCR
+do MPRJ e aguardar o resultado do processamento.
+
+A API IntelliDoc processa documentos de forma assíncrona:
+1. POST /documents/files → submete documento, retorna document_id
+2. GET /documents/{id} → consulta status/resultado
+
+Example:
+    Uso básico com variável de ambiente:
+
+    >>> from nia_etl_utils import executar_ocr, OcrError
+    >>>
+    >>> try:
+    ...     resultado = executar_ocr(blob_bytes, url_base="INTELLIDOC_URL")
+    ...     print(resultado["full_text"])
+    ... except OcrError as e:
+    ...     logger.error(f"Falha no OCR: {e}")
+"""
+
+import time
+
+import requests
+from loguru import logger
+
+from .env_config import obter_variavel_env
+from .exceptions import (
+    OcrProcessamentoError,
+    OcrSubmissaoError,
+    OcrTimeoutError,
+)
+
+# =============================================================================
+# CONSTANTES
+# =============================================================================
+
+MAGIC_BYTES: dict[bytes, str] = {
+    b"%PDF": ".pdf",
+    b"\xff\xd8\xff": ".jpg",
+    b"\x89PNG\r\n\x1a\n": ".png",
+    b"GIF87a": ".gif",
+    b"GIF89a": ".gif",
+    b"BM": ".bmp",
+    b"II*\x00": ".tiff",
+    b"MM\x00*": ".tiff",
+}
+
+PATH_ENVIO = "/documents/files"
+PATH_CONSULTA = "/documents"
+
+
+# =============================================================================
+# FUNÇÕES AUXILIARES
+# =============================================================================
+
+
+def _detectar_extensao(conteudo: bytes) -> str:
+    """Detecta a extensão do arquivo baseado nos magic bytes.
+
+    Args:
+        conteudo: Bytes do arquivo.
+
+    Returns:
+        Extensão detectada (ex: '.pdf', '.jpg') ou '.bin' se não reconhecido.
+    """
+    conteudo_inicio = conteudo[:32].lstrip()
+
+    for magic, extensao in MAGIC_BYTES.items():
+        if conteudo_inicio.startswith(magic):
+            return extensao
+
+    return ".bin"
+
+
+def _normalizar_para_bytes(blob: object) -> bytes:
+    """Converte diferentes tipos de BLOB para bytes.
+
+    Suporta LOBs do cx_Oracle/oracledb, memoryview e bytes/bytearray.
+
+    Args:
+        blob: Objeto contendo dados binários.
+
+    Returns:
+        Conteúdo em bytes.
+
+    Raises:
+        TypeError: Se o tipo do blob não for suportado.
+    """
+    if hasattr(blob, "read"):
+        return blob.read()  # type: ignore[union-attr]
+
+    if isinstance(blob, memoryview):
+        return blob.tobytes()
+
+    if isinstance(blob, (bytes, bytearray)):
+        return bytes(blob)
+
+    raise TypeError(f"Tipo de blob não suportado: {type(blob)}")
+
+
+def _resolver_url_base(url_base: str) -> str:
+    """Resolve URL base a partir de valor direto ou nome de variável de ambiente.
+
+    Args:
+        url_base: URL direta (começa com http) ou nome de variável de ambiente.
+
+    Returns:
+        URL base resolvida, sem barra final.
+    """
+    if url_base.startswith(("http://", "https://")):  # noqa
+        url = url_base
+    else:
+        url = obter_variavel_env(url_base)
+
+    return url.rstrip("/")
+
+
+def _submeter_documento(
+    url_base: str,
+    conteudo: bytes,
+    extensao: str,
+    max_tentativas: int,
+    intervalo_retry: int,
+) -> str:
+    """Submete documento para processamento OCR.
+
+    Args:
+        url_base: URL base da API.
+        conteudo: Bytes do documento.
+        extensao: Extensão do arquivo.
+        max_tentativas: Número máximo de tentativas.
+        intervalo_retry: Segundos entre tentativas.
+
+    Returns:
+        document_id retornado pela API.
+
+    Raises:
+        OcrSubmissaoError: Se todas as tentativas falharem.
+    """
+    url = f"{url_base}{PATH_ENVIO}"
+    nome_arquivo = f"documento{extensao}"
+    files = {"files": (nome_arquivo, conteudo)}
+
+    ultima_excecao: Exception | None = None
+
+    for tentativa in range(1, max_tentativas + 1):
+        try:
+            logger.debug(
+                f"Tentativa {tentativa}/{max_tentativas} - "
+                f"Submetendo OCR (extensao={extensao}, tamanho={len(conteudo)} bytes)"
+            )
+
+            resp = requests.post(url, files=files, timeout=120)
+
+            if resp.ok:
+                resultado = resp.json()
+                document_id = resultado["documents"][0]["document_id"]
+                logger.info(f"Documento submetido com sucesso. document_id={document_id}")
+                return document_id
+
+            try:
+                detalhe = resp.json()
+            except Exception:
+                detalhe = resp.text
+
+            ultima_excecao = OcrSubmissaoError(
+                f"API retornou status {resp.status_code}",
+                details={"status": resp.status_code, "detalhe": detalhe},
+            )
+            logger.warning(f"Tentativa {tentativa} falhou: status={resp.status_code}")
+
+        except requests.RequestException as e:
+            ultima_excecao = e
+            logger.warning(f"Tentativa {tentativa} falhou com exceção: {e}")
+
+        if tentativa < max_tentativas:
+            logger.info(f"Aguardando {intervalo_retry}s antes da próxima tentativa...")
+            time.sleep(intervalo_retry)
+
+    raise OcrSubmissaoError(
+        f"Submissão OCR falhou após {max_tentativas} tentativas",
+        details={"tentativas": max_tentativas, "ultimo_erro": str(ultima_excecao)},
+    )
+
+
+def _aguardar_resultado(
+    url_base: str,
+    document_id: str,
+    timeout_polling: int,
+    intervalo_polling: int,
+) -> dict:
+    """Aguarda resultado do processamento OCR via polling.
+
+    Args:
+        url_base: URL base da API.
+        document_id: ID do documento retornado na submissão.
+        timeout_polling: Tempo máximo de espera em segundos.
+        intervalo_polling: Intervalo entre consultas em segundos.
+
+    Returns:
+        Dicionário com resultado completo da API (campo 'result').
+
+    Raises:
+        OcrProcessamentoError: Se a API retornar FAILURE ou REVOKED.
+        OcrTimeoutError: Se o timeout for atingido.
+    """
+    url = f"{url_base}{PATH_CONSULTA}/{document_id}"
+    inicio = time.time()
+    ultimo_status = "UNKNOWN"
+
+    while True:
+        tempo_decorrido = time.time() - inicio
+
+        if tempo_decorrido >= timeout_polling:
+            raise OcrTimeoutError(
+                f"Timeout após {timeout_polling}s aguardando OCR",
+                details={"document_id": document_id, "ultimo_status": ultimo_status},
+            )
+
+        try:
+            resp = requests.get(url, timeout=30)
+
+            if resp.ok:
+                resultado = resp.json()
+                ultimo_status = resultado.get("status", "UNKNOWN")
+
+                if ultimo_status == "SUCCESS":
+                    logger.info(f"OCR concluído para document_id={document_id}")
+                    return resultado.get("result", {})
+
+                if ultimo_status in ("FAILURE", "REVOKED"):
+                    erro = resultado.get("error") or resultado.get("message") or "Erro desconhecido"
+                    raise OcrProcessamentoError(
+                        f"OCR falhou: {erro}",
+                        details={"document_id": document_id, "status": ultimo_status, "erro_api": erro},
+                    )
+
+                logger.debug(f"document_id={document_id} status={ultimo_status}, aguardando...")
+
+        except requests.RequestException as e:
+            logger.warning(f"Erro ao consultar status: {e}")
+
+        time.sleep(intervalo_polling)
+
+
+# =============================================================================
+# FUNÇÃO PRINCIPAL
+# =============================================================================
+
+
+def executar_ocr(
+    conteudo: bytes | object,
+    url_base: str = "INTELLIDOC_URL",
+    timeout_polling: int = 300,
+    max_tentativas: int = 3,
+    intervalo_retry: int = 5,
+    intervalo_polling: int = 1,
+) -> dict:
+    """Executa OCR em documento via API IntelliDoc.
+
+    Submete o documento para processamento e aguarda o resultado.
+    A extensão do arquivo é detectada automaticamente pelos magic bytes.
+
+    Args:
+        conteudo: Bytes do documento ou objeto com método read() (LOB Oracle).
+        url_base: URL da API ou nome da variável de ambiente.
+            Se começar com 'http://' ou 'https://', usa como URL direta.
+            Caso contrário, busca no .env. Default: "INTELLIDOC_URL".
+        timeout_polling: Tempo máximo em segundos para aguardar resultado.
+            Default: 300 (5 minutos).
+        max_tentativas: Número de tentativas para submissão. Default: 3.
+        intervalo_retry: Segundos entre tentativas de submissão. Default: 5.
+        intervalo_polling: Segundos entre consultas de status. Default: 1.
+
+    Returns:
+        Dicionário com resultado completo da API, contendo:
+            - document_id: ID do documento
+            - full_text: Texto extraído completo
+            - mime_type: Tipo MIME detectado
+            - overall_quality: Qualidade geral do OCR (0-1)
+            - total_pages: Número de páginas
+            - processing_time_ms: Tempo de processamento
+            - pages: Lista com detalhes de cada página
+            - metadata: Metadados adicionais
+
+    Raises:
+        OcrSubmissaoError: Se falhar ao submeter documento.
+        OcrProcessamentoError: Se a API retornar erro no processamento.
+        OcrTimeoutError: Se timeout for atingido aguardando resultado.
+        TypeError: Se o tipo do conteúdo não for suportado.
+
+    Examples:
+        Uso básico:
+
+        >>> resultado = executar_ocr(blob_bytes, url_base="INTELLIDOC_URL")
+        >>> print(resultado["full_text"])
+
+        Uso com URL direta:
+
+        >>> resultado = executar_ocr(
+        ...     conteudo=pdf_bytes,
+        ...     url_base="http://intellidoc.mprj.mp.br",
+        ...     timeout_polling=600,
+        ... )
+        >>> print(f"Qualidade: {resultado['overall_quality']}")
+    """
+    conteudo_bytes = _normalizar_para_bytes(conteudo)
+    extensao = _detectar_extensao(conteudo_bytes)
+    url = _resolver_url_base(url_base)
+
+    logger.info(f"Iniciando OCR (tamanho={len(conteudo_bytes)} bytes, extensao={extensao})")
+
+    document_id = _submeter_documento(
+        url_base=url,
+        conteudo=conteudo_bytes,
+        extensao=extensao,
+        max_tentativas=max_tentativas,
+        intervalo_retry=intervalo_retry,
+    )
+
+    return _aguardar_resultado(
+        url_base=url,
+        document_id=document_id,
+        timeout_polling=timeout_polling,
+        intervalo_polling=intervalo_polling,
+    )