283 lines
11 KiB
Python
283 lines
11 KiB
Python
"""Tareas ARQ para procesamiento paralelo de documentos.
|
|
|
|
Cada tarea recibe la ruta de un archivo temporal en el volumen compartido
|
|
`/tmp/ips_queue/`, lo procesa y actualiza la tarea en la base de datos.
|
|
El archivo se elimina al finalizar (éxito o error).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
from typing import Any
|
|
|
|
from sqlalchemy.orm.attributes import flag_modified
|
|
from infrastructure.database import ExtractionRow, Task, session_scope, utcnow
|
|
from infrastructure.serialization import flatten_extraction_result
|
|
from api.desktop_ingest import (
|
|
_finalize_radicado_if_ready,
|
|
_paths_for_type,
|
|
_service_for_model,
|
|
mark_task_completed_if_done,
|
|
)
|
|
from domain.profiles import get_profile
|
|
from domain.prompts import PromptBuilder
|
|
from domain.normalizer import normalize_to_schema
|
|
from domain.rate_limiter import get_rate_limiter
|
|
from domain.config import settings
|
|
from domain.usage_tracking import record_api_usage
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _service_for_request(
|
|
ctx: dict[str, Any], request: dict[str, Any]
|
|
) -> Any:
|
|
"""Resuelve y reutiliza el servicio correspondiente al modelo solicitado."""
|
|
base_service = ctx["service"]
|
|
requested_model = str(request.get("model") or "").strip() or None
|
|
if not requested_model:
|
|
return base_service
|
|
|
|
current_model = getattr(getattr(base_service, "extractor", None), "model", None)
|
|
if current_model == requested_model:
|
|
return base_service
|
|
|
|
services_by_model = ctx.setdefault("services_by_model", {})
|
|
service = services_by_model.get(requested_model)
|
|
if service is None:
|
|
service = _service_for_model(base_service, requested_model)
|
|
services_by_model[requested_model] = service
|
|
return service
|
|
|
|
|
|
async def process_document(
|
|
ctx: dict[str, Any],
|
|
*,
|
|
task_id: str,
|
|
batch_id: str,
|
|
queue_path: str,
|
|
radicado: str,
|
|
filename: str,
|
|
doc_type: str | None,
|
|
instruction: str | None,
|
|
tiff_page_range: list[int] | None,
|
|
) -> dict[str, Any]:
|
|
"""Extrae datos de un documento y los persiste en la BD.
|
|
|
|
Args:
|
|
ctx: Contexto ARQ con 'service' (InvoiceProcessingService).
|
|
task_id: ID de la tarea parent.
|
|
batch_id: ID del lote de extracción.
|
|
queue_path: Ruta absoluta al archivo en el volumen compartido.
|
|
radicado: Radicado al que pertenece el documento.
|
|
filename: Nombre original del archivo.
|
|
doc_type: Tipo documental (None en modo carpeta general).
|
|
instruction: Instrucción de extracción personalizada.
|
|
tiff_page_range: Rango de páginas TIFF [inicio, fin].
|
|
"""
|
|
rate_limiter = get_rate_limiter()
|
|
tiff_range = tuple(tiff_page_range) if tiff_page_range else None
|
|
|
|
logger.info("Procesando %s / %s (task=%s)", radicado, filename, task_id)
|
|
|
|
try:
|
|
# Obtener config de la tarea desde BD
|
|
def _load_task_config() -> dict[str, Any] | None:
|
|
with session_scope() as session:
|
|
task = session.get(Task, task_id)
|
|
if task is None:
|
|
return None
|
|
return dict(task.request or {})
|
|
|
|
request = await asyncio.to_thread(_load_task_config)
|
|
if request is None:
|
|
logger.warning("Tarea %s no encontrada — saltando %s", task_id, filename)
|
|
return {"skipped": True, "reason": "task_not_found"}
|
|
|
|
service = _service_for_request(ctx, request)
|
|
effective_model = getattr(getattr(service, "extractor", None), "model", "?")
|
|
logger.info(
|
|
"Modelo resuelto para task=%s: %s",
|
|
task_id,
|
|
effective_model,
|
|
)
|
|
|
|
fields: list[str] = list(request.get("fields") or [])
|
|
alias_map: dict[str, str] = dict(request.get("alias_map") or {})
|
|
structure_mode: str = str(request.get("structure_mode") or "radicados")
|
|
doc_types: list = list(request.get("doc_types") or [])
|
|
prof = get_profile("invoice")
|
|
|
|
selected_paths = _paths_for_type(
|
|
fields,
|
|
doc_type if structure_mode == "radicados" and doc_types else None,
|
|
)
|
|
|
|
if doc_type and structure_mode == "radicados" and doc_types:
|
|
try:
|
|
prompt = PromptBuilder.build_for_type(
|
|
doc_type, selected_paths, {}, profile_name=prof.get("name", "invoice")
|
|
)
|
|
except Exception:
|
|
prompt = PromptBuilder.build(
|
|
selected_paths, {}, profile_name=prof.get("name", "invoice")
|
|
)
|
|
else:
|
|
prompt = instruction or PromptBuilder.build(
|
|
selected_paths, {}, profile_name=prof.get("name", "invoice")
|
|
)
|
|
|
|
# Esperar turno en el rate limiter antes de llamar a OpenAI
|
|
estimated_tokens = max(len(selected_paths) * 80, 2000)
|
|
await rate_limiter.acquire(estimated_tokens=estimated_tokens)
|
|
|
|
# Extracción + normalización (CPU + IO bound, en thread pool)
|
|
def _extract_and_normalize() -> tuple[dict, dict, dict | None, dict | None]:
|
|
data, extraction_usage = service.extract_full_from_pdf(
|
|
queue_path,
|
|
instruction=prompt,
|
|
tiff_page_range=tiff_range,
|
|
)
|
|
raw_data: dict = data if isinstance(data, dict) else {"data": data}
|
|
if isinstance(data, dict):
|
|
data["radicado"] = radicado
|
|
else:
|
|
data = {"radicado": radicado, "data": data}
|
|
|
|
norm = normalize_to_schema(
|
|
data,
|
|
paths=selected_paths,
|
|
profile_name=prof.get("name", "invoice"),
|
|
)
|
|
normalized_obj = data
|
|
normalizer_usage = None
|
|
if isinstance(norm, dict):
|
|
normalized_obj = norm.get("normalized") or data
|
|
normalizer_usage = norm.get("_usage")
|
|
return raw_data, normalized_obj, extraction_usage, normalizer_usage
|
|
|
|
_raw, normalized_obj, extraction_usage, normalizer_usage = await asyncio.to_thread(_extract_and_normalize)
|
|
|
|
# Persistir resultado en BD
|
|
def _persist() -> dict[str, Any]:
|
|
with session_scope() as session:
|
|
# Lock evita lost-update en result.processed con N workers concurrentes.
|
|
task = session.get(Task, task_id, with_for_update=True)
|
|
if task is None:
|
|
return {"skipped": True, "reason": "task_not_found"}
|
|
|
|
result = dict(task.result or {})
|
|
result = record_api_usage(
|
|
result, phase="extract", model=str(effective_model), usage=extraction_usage,
|
|
radicado=radicado, filename=filename,
|
|
)
|
|
result = record_api_usage(
|
|
result, phase="normalize", model=str(settings.api.normalizer_model), usage=normalizer_usage,
|
|
radicado=radicado, filename=filename,
|
|
)
|
|
task.result = result
|
|
|
|
if structure_mode == "radicados" and doc_types and doc_type:
|
|
partials = dict(result.get("partials") or {})
|
|
radicado_partials = dict(partials.get(radicado) or {})
|
|
radicado_partials[doc_type] = normalized_obj or {}
|
|
partials[radicado] = radicado_partials
|
|
result["partials"] = partials
|
|
task.result = result
|
|
session.flush()
|
|
finalized = _finalize_radicado_if_ready(
|
|
session,
|
|
task=task,
|
|
batch_id=batch_id,
|
|
radicado=radicado,
|
|
request=request,
|
|
fields=fields,
|
|
alias_map=alias_map,
|
|
)
|
|
mark_task_completed_if_done(session, task, batch_id)
|
|
return {
|
|
"radicado": radicado,
|
|
"filename": filename,
|
|
"doc_type": doc_type,
|
|
"finalized": finalized,
|
|
"processed": int((task.result or {}).get("processed") or 0),
|
|
}
|
|
|
|
fused = {"radicado": radicado, "_source_pdf": filename}
|
|
if doc_type:
|
|
fused[doc_type] = normalized_obj or {}
|
|
else:
|
|
fused.update(
|
|
normalized_obj if isinstance(normalized_obj, dict) else {"data": normalized_obj}
|
|
)
|
|
|
|
row_index = int(result.get("row_count") or 0)
|
|
for row in flatten_extraction_result(fused, fields, alias_map):
|
|
row_index += 1
|
|
session.add(
|
|
ExtractionRow(
|
|
batch_id=batch_id,
|
|
radicado=radicado,
|
|
source_pdf=filename,
|
|
row_index=row_index,
|
|
row_json=row,
|
|
raw_json=fused,
|
|
)
|
|
)
|
|
|
|
result["row_count"] = row_index
|
|
result["processed"] = int(result.get("processed") or 0) + 1
|
|
task.result = result
|
|
flag_modified(task, "result")
|
|
task.updated_at = utcnow()
|
|
session.flush()
|
|
mark_task_completed_if_done(session, task, batch_id)
|
|
return {
|
|
"radicado": radicado,
|
|
"filename": filename,
|
|
"doc_type": doc_type,
|
|
"finalized": True,
|
|
"processed": int(result.get("processed") or 0),
|
|
}
|
|
|
|
payload = await asyncio.to_thread(_persist)
|
|
logger.info("Completado %s / %s — payload=%s", radicado, filename, payload)
|
|
return payload
|
|
|
|
except Exception as exc:
|
|
logger.exception("Error procesando %s / %s: %s", radicado, filename, exc)
|
|
# Registrar error en la tarea sin fallar el worker
|
|
def _record_error() -> None:
|
|
try:
|
|
with session_scope() as session:
|
|
task = session.get(Task, task_id, with_for_update=True)
|
|
if task is None:
|
|
return
|
|
result = dict(task.result or {})
|
|
errors = list(result.get("errors") or [])
|
|
errors.append({
|
|
"radicado": radicado,
|
|
"filename": filename,
|
|
"error": str(exc),
|
|
})
|
|
result["errors"] = errors
|
|
result["processed"] = int(result.get("processed") or 0) + 1
|
|
task.result = result
|
|
flag_modified(task, "result")
|
|
task.updated_at = utcnow()
|
|
mark_task_completed_if_done(session, task, batch_id)
|
|
except Exception as inner:
|
|
logger.error("No se pudo registrar el error en BD: %s", inner)
|
|
|
|
await asyncio.to_thread(_record_error)
|
|
return {"error": str(exc), "radicado": radicado, "filename": filename}
|
|
|
|
finally:
|
|
# Eliminar archivo temporal del volumen compartido
|
|
try:
|
|
if os.path.exists(queue_path):
|
|
os.unlink(queue_path)
|
|
except OSError as e:
|
|
logger.warning("No se pudo eliminar %s: %s", queue_path, e)
|