feat: Add .env.example configuration file and integrate Redis service into Docker Compose for improved task management. Update desktop application to support document indexing and extraction workflows, enhancing user experience with a new extraction wizard and stepper components. Revise README and documentation for clarity on setup and usage.
This commit is contained in:
@@ -17,6 +17,7 @@ COPY domain /app/domain
|
||||
COPY core /app/core
|
||||
COPY infrastructure /app/infrastructure
|
||||
COPY api /app/api
|
||||
COPY worker /app/worker
|
||||
|
||||
EXPOSE 8001
|
||||
CMD ["uvicorn", "api.main:app", "--host", "0.0.0.0", "--port", "8001"]
|
||||
|
||||
@@ -37,7 +37,9 @@ en `dependencies/poppler-24.08.0/Library/bin`.
|
||||
- `GET /config`
|
||||
- `GET /models`
|
||||
- `POST /schemas`
|
||||
- `POST /schemas/upload`
|
||||
- `POST /estimate`
|
||||
- `POST /batches`
|
||||
- `GET /batches/{batch_id}`
|
||||
- `GET /batches/{batch_id}/rows`
|
||||
- `POST /tasks/extraction`
|
||||
- `POST /tasks/{task_id}/ingest`
|
||||
- `GET /tasks`
|
||||
|
||||
@@ -0,0 +1,235 @@
|
||||
"""Procesamiento de ingestas individuales desde la app desktop."""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from domain.normalizer import normalize_to_schema
|
||||
from domain.profiles import get_profile
|
||||
from domain.prompts import PromptBuilder
|
||||
from domain.service import InvoiceProcessingService
|
||||
from infrastructure.database import ExtractionRow, Task, utcnow
|
||||
from infrastructure.serialization import flatten_extraction_result
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _paths_for_type(fields: list[str], doc_type: str | None) -> list[str]:
|
||||
if not doc_type:
|
||||
return list(fields or [])
|
||||
prefix = f"{doc_type}/"
|
||||
return [path.split("/", 1)[1] for path in fields if path.startswith(prefix)]
|
||||
|
||||
|
||||
def _expected_doc_types_for_radicado(request: dict[str, Any], radicado: str) -> set[str]:
|
||||
expected: set[str] = set()
|
||||
for item in request.get("documents") or []:
|
||||
if str(item.get("radicado") or "") != radicado:
|
||||
continue
|
||||
doc_type = str(item.get("doc_type") or item.get("docType") or "").strip()
|
||||
if doc_type:
|
||||
expected.add(doc_type)
|
||||
return expected
|
||||
|
||||
|
||||
def _finalize_radicado_if_ready(
|
||||
session: Session,
|
||||
*,
|
||||
task: Task,
|
||||
batch_id: str,
|
||||
radicado: str,
|
||||
request: dict[str, Any],
|
||||
fields: list[str],
|
||||
alias_map: dict[str, str],
|
||||
) -> bool:
|
||||
partials = dict((task.result or {}).get("partials") or {})
|
||||
radicado_partials = dict(partials.get(radicado) or {})
|
||||
expected = _expected_doc_types_for_radicado(request, radicado)
|
||||
if expected and not expected.issubset(set(radicado_partials.keys())):
|
||||
return False
|
||||
|
||||
fused: dict[str, Any] = {"radicado": radicado, **radicado_partials}
|
||||
source_pdf = None
|
||||
for item in request.get("documents") or []:
|
||||
if str(item.get("radicado") or "") == radicado:
|
||||
source_pdf = item.get("filename")
|
||||
break
|
||||
if source_pdf:
|
||||
fused["_source_pdf"] = source_pdf
|
||||
|
||||
row_index = int((task.result or {}).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=source_pdf,
|
||||
row_index=row_index,
|
||||
row_json=row,
|
||||
raw_json=fused,
|
||||
)
|
||||
)
|
||||
|
||||
partials.pop(radicado, None)
|
||||
result = dict(task.result or {})
|
||||
result["partials"] = partials
|
||||
result["row_count"] = row_index
|
||||
result["processed"] = int(result.get("processed") or 0) + 1
|
||||
task.result = result
|
||||
task.updated_at = utcnow()
|
||||
return True
|
||||
|
||||
|
||||
def _service_for_model(base_service: InvoiceProcessingService, model: str | None) -> InvoiceProcessingService:
|
||||
"""Devuelve un servicio con el modelo solicitado; reutiliza base_service si coincide."""
|
||||
if not model:
|
||||
return base_service
|
||||
current_model = getattr(base_service.extractor, "model", None)
|
||||
if current_model == model:
|
||||
return base_service
|
||||
from domain.strategies import OpenAICompatibleExtractor
|
||||
from domain.batch_processor import BatchProcessor
|
||||
from domain.estimator import Estimator
|
||||
extractor = OpenAICompatibleExtractor(model=model)
|
||||
return InvoiceProcessingService(
|
||||
extractor,
|
||||
batch_processor_factory=lambda s: BatchProcessor(s),
|
||||
estimator_factory=lambda s: Estimator(s),
|
||||
)
|
||||
|
||||
|
||||
def process_desktop_ingest(
|
||||
service: InvoiceProcessingService,
|
||||
session: Session,
|
||||
*,
|
||||
task: Task,
|
||||
batch_id: str,
|
||||
temp_path: str,
|
||||
radicado: str,
|
||||
filename: str,
|
||||
doc_type: str | None,
|
||||
instruction: str | None,
|
||||
tiff_page_range: tuple[int, int] | None,
|
||||
model: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
service = _service_for_model(service, model)
|
||||
request = dict(task.request or {})
|
||||
fields = list(request.get("fields") or [])
|
||||
alias_map = dict(request.get("alias_map") or {})
|
||||
structure_mode = str(request.get("structure_mode") or "radicados")
|
||||
doc_types = 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"))
|
||||
|
||||
data, _usage = service.extract_full_from_pdf(
|
||||
temp_path,
|
||||
instruction=prompt,
|
||||
tiff_page_range=tiff_page_range,
|
||||
)
|
||||
raw_data = 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
|
||||
if isinstance(norm, dict):
|
||||
normalized_obj = norm.get("normalized") or data
|
||||
|
||||
result = dict(task.result or {})
|
||||
errors = list(result.get("errors") or [])
|
||||
|
||||
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,
|
||||
)
|
||||
return {
|
||||
"radicado": radicado,
|
||||
"filename": filename,
|
||||
"doc_type": doc_type,
|
||||
"finalized": finalized,
|
||||
"processed": int((task.result or {}).get("processed") or 0),
|
||||
"total": int((task.result or {}).get("total") 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
|
||||
task.updated_at = utcnow()
|
||||
session.flush()
|
||||
|
||||
return {
|
||||
"radicado": radicado,
|
||||
"filename": filename,
|
||||
"doc_type": doc_type,
|
||||
"finalized": True,
|
||||
"processed": int(result.get("processed") or 0),
|
||||
"total": int(result.get("total") or 0),
|
||||
}
|
||||
|
||||
|
||||
def mark_task_completed_if_done(session: Session, task: Task, batch_id: str) -> None:
|
||||
result = dict(task.result or {})
|
||||
processed = int(result.get("processed") or 0)
|
||||
total = int(result.get("total") or 0)
|
||||
partials = dict(result.get("partials") or {})
|
||||
if processed < total or partials:
|
||||
return
|
||||
|
||||
from infrastructure.database import ExtractionBatch
|
||||
|
||||
task.status = "completed"
|
||||
task.finished_at = utcnow()
|
||||
task.updated_at = utcnow()
|
||||
batch = session.get(ExtractionBatch, batch_id)
|
||||
if batch is not None:
|
||||
batch.status = "completed"
|
||||
batch.updated_at = utcnow()
|
||||
+492
-33
@@ -1,13 +1,26 @@
|
||||
"""API REST del microservicio de extraccion."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import uuid
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
from fastapi import BackgroundTasks, FastAPI, HTTPException
|
||||
from pydantic import BaseModel, Field
|
||||
try:
|
||||
from arq import create_pool as _arq_create_pool
|
||||
from arq.connections import RedisSettings as _RedisSettings
|
||||
_ARQ_AVAILABLE = True
|
||||
except ImportError:
|
||||
_arq_create_pool = None # type: ignore[assignment]
|
||||
_RedisSettings = None # type: ignore[assignment]
|
||||
_ARQ_AVAILABLE = False
|
||||
|
||||
from fastapi import BackgroundTasks, FastAPI, File, Form, HTTPException, UploadFile
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
from sqlalchemy import select
|
||||
|
||||
from api.desktop_ingest import mark_task_completed_if_done, process_desktop_ingest
|
||||
|
||||
from domain.batch_processor import BatchProcessor
|
||||
from domain.config import extractor_llm_config, normalizer_llm_config, settings
|
||||
from domain.config import get_allowed_document_types
|
||||
@@ -23,6 +36,7 @@ from infrastructure.database import (
|
||||
utcnow,
|
||||
)
|
||||
from infrastructure.serialization import flatten_extraction_result
|
||||
from infrastructure.uploads import temp_upload_path
|
||||
|
||||
|
||||
def build_processing_service() -> InvoiceProcessingService:
|
||||
@@ -42,10 +56,43 @@ def build_processing_service() -> InvoiceProcessingService:
|
||||
processing_service = build_processing_service()
|
||||
|
||||
|
||||
def _redis_settings_from_url(url: str):
|
||||
url = url.removeprefix("redis://")
|
||||
host_port, *db_parts = url.split("/")
|
||||
host, *port_parts = host_port.split(":")
|
||||
port = int(port_parts[0]) if port_parts else 6379
|
||||
db = int(db_parts[0]) if db_parts else 0
|
||||
return _RedisSettings(host=host, port=port, database=db)
|
||||
|
||||
|
||||
# Pool ARQ global — compartido por todos los requests
|
||||
_arq_pool = None
|
||||
|
||||
|
||||
async def get_arq_pool():
|
||||
global _arq_pool
|
||||
if not _ARQ_AVAILABLE:
|
||||
raise RuntimeError("arq no instalado; ASYNC_QUEUE requiere 'pip install arq'.")
|
||||
if _arq_pool is None:
|
||||
_arq_pool = await _arq_create_pool(_redis_settings_from_url(settings.redis_url))
|
||||
return _arq_pool
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
init_db()
|
||||
if settings.processing.async_queue:
|
||||
os.makedirs(settings.queue_dir, exist_ok=True)
|
||||
try:
|
||||
await get_arq_pool()
|
||||
except Exception as exc:
|
||||
import logging as _logging
|
||||
_logging.getLogger(__name__).warning(
|
||||
"No se pudo conectar a Redis (%s). Usando modo sincrono como fallback.", exc
|
||||
)
|
||||
yield
|
||||
if _arq_pool is not None:
|
||||
await _arq_pool.aclose()
|
||||
|
||||
|
||||
app = FastAPI(
|
||||
@@ -69,8 +116,15 @@ class SchemaRequest(BaseModel):
|
||||
structure_mode: str = "radicados"
|
||||
|
||||
|
||||
class DesktopDocument(BaseModel):
|
||||
radicado: str
|
||||
filename: str
|
||||
doc_type: str | None = None
|
||||
|
||||
|
||||
class ExtractionRequest(BaseModel):
|
||||
root_path: str
|
||||
source: str = "legacy"
|
||||
root_path: str = ""
|
||||
keywords: list[str] | None = None
|
||||
doc_types: list[DocumentType] | None = None
|
||||
fields: list[str] = Field(default_factory=list)
|
||||
@@ -80,6 +134,18 @@ class ExtractionRequest(BaseModel):
|
||||
structure_mode: str = "radicados"
|
||||
retry_on_empty: bool = False
|
||||
radicado_filter: list[str] | None = None
|
||||
documents: list[DesktopDocument] | None = None
|
||||
document_count: int | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_source(self) -> "ExtractionRequest":
|
||||
if self.source == "desktop":
|
||||
if not self.documents:
|
||||
raise ValueError("Envia documents cuando source=desktop.")
|
||||
return self
|
||||
if not self.root_path:
|
||||
raise ValueError("Envia root_path para extraccion legacy.")
|
||||
return self
|
||||
|
||||
|
||||
class EstimateRequest(BaseModel):
|
||||
@@ -309,29 +375,34 @@ def list_models() -> dict[str, Any]:
|
||||
@app.post("/schemas")
|
||||
def build_schema(request: SchemaRequest) -> dict[str, Any]:
|
||||
"""Descubre campos disponibles para que el cliente API elija `fields`."""
|
||||
if request.document_path:
|
||||
schema, sample = processing_service.build_schema_from_first(
|
||||
request.document_path,
|
||||
instruction=request.instruction,
|
||||
tiff_page_range=request.tiff_page_range,
|
||||
)
|
||||
return {"schema": schema, "sample": sample}
|
||||
|
||||
if request.root_path and request.doc_types:
|
||||
schema_by_type, sample_by_type, first_radicado = (
|
||||
processing_service.build_schema_from_first_for_types(
|
||||
request.root_path,
|
||||
[dt.model_dump() for dt in request.doc_types],
|
||||
try:
|
||||
if request.document_path:
|
||||
schema, sample = processing_service.build_schema_from_first(
|
||||
request.document_path,
|
||||
instruction=request.instruction,
|
||||
tiff_page_range=request.tiff_page_range,
|
||||
structure_mode=request.structure_mode,
|
||||
)
|
||||
)
|
||||
return {
|
||||
"schema_by_type": schema_by_type,
|
||||
"sample_by_type": sample_by_type,
|
||||
"first_radicado": first_radicado,
|
||||
}
|
||||
return {"schema": schema, "sample": sample}
|
||||
|
||||
if request.root_path and request.doc_types:
|
||||
schema_by_type, sample_by_type, first_radicado = (
|
||||
processing_service.build_schema_from_first_for_types(
|
||||
request.root_path,
|
||||
[dt.model_dump() for dt in request.doc_types],
|
||||
instruction=request.instruction,
|
||||
tiff_page_range=request.tiff_page_range,
|
||||
structure_mode=request.structure_mode,
|
||||
)
|
||||
)
|
||||
return {
|
||||
"schema_by_type": schema_by_type,
|
||||
"sample_by_type": sample_by_type,
|
||||
"first_radicado": first_radicado,
|
||||
}
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
except FileNotFoundError as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
||||
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
@@ -339,18 +410,92 @@ def build_schema(request: SchemaRequest) -> dict[str, Any]:
|
||||
)
|
||||
|
||||
|
||||
@app.post("/schemas/upload")
|
||||
def build_schema_upload(
|
||||
file: UploadFile = File(...),
|
||||
structure_mode: str = Form("radicados"),
|
||||
doc_type_name: str | None = Form(None),
|
||||
instruction: str | None = Form(None),
|
||||
tiff_start: int | None = Form(None),
|
||||
tiff_end: int | None = Form(None),
|
||||
) -> dict[str, Any]:
|
||||
"""Descubre campos a partir de un documento subido desde la app desktop."""
|
||||
tiff_page_range = None
|
||||
if tiff_start is not None and tiff_end is not None and not (tiff_start == 0 and tiff_end == 0):
|
||||
tiff_page_range = (tiff_start, tiff_end)
|
||||
|
||||
try:
|
||||
with temp_upload_path(file) as temp_path:
|
||||
schema, sample = processing_service.build_schema_from_first(
|
||||
temp_path,
|
||||
instruction=instruction,
|
||||
tiff_page_range=tiff_page_range,
|
||||
)
|
||||
if doc_type_name:
|
||||
return {
|
||||
"schema_by_type": {doc_type_name: schema},
|
||||
"sample_by_type": {doc_type_name: sample},
|
||||
"structure_mode": structure_mode,
|
||||
}
|
||||
return {"schema": schema, "sample": sample, "structure_mode": structure_mode}
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
except FileNotFoundError as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
||||
|
||||
|
||||
@app.post("/estimate")
|
||||
def estimate(request: EstimateRequest) -> dict[str, Any]:
|
||||
return processing_service.estimate_tokens_and_cost(
|
||||
request.root_path,
|
||||
request.keywords,
|
||||
request.schema_map,
|
||||
request.selected_fields,
|
||||
doc_types=[dt.model_dump() for dt in request.doc_types or []] or None,
|
||||
schema_by_type=request.schema_by_type,
|
||||
selected_paths=request.selected_paths,
|
||||
structure_mode=request.structure_mode,
|
||||
)
|
||||
try:
|
||||
return processing_service.estimate_tokens_and_cost(
|
||||
request.root_path,
|
||||
request.keywords,
|
||||
request.schema_map,
|
||||
request.selected_fields,
|
||||
doc_types=[dt.model_dump() for dt in request.doc_types or []] or None,
|
||||
schema_by_type=request.schema_by_type,
|
||||
selected_paths=request.selected_paths,
|
||||
structure_mode=request.structure_mode,
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
|
||||
|
||||
@app.post("/estimate/upload")
|
||||
def estimate_upload(
|
||||
file: UploadFile = File(...),
|
||||
doc_count: int = Form(1),
|
||||
selected_paths: str | None = Form(None),
|
||||
structure_mode: str = Form("radicados"),
|
||||
tiff_start: int | None = Form(None),
|
||||
tiff_end: int | None = Form(None),
|
||||
model: str | None = Form(None),
|
||||
) -> dict[str, Any]:
|
||||
"""Estima tokens y costos a partir de un documento subido desde la app desktop."""
|
||||
import json as _json
|
||||
|
||||
tiff_page_range = None
|
||||
if tiff_start is not None and tiff_end is not None and not (tiff_start == 0 and tiff_end == 0):
|
||||
tiff_page_range = (tiff_start, tiff_end)
|
||||
|
||||
paths_list: list[str] = []
|
||||
if selected_paths:
|
||||
try:
|
||||
paths_list = _json.loads(selected_paths)
|
||||
except Exception:
|
||||
paths_list = []
|
||||
|
||||
try:
|
||||
with temp_upload_path(file) as temp_path:
|
||||
return processing_service.estimator.estimate_from_file(
|
||||
temp_path,
|
||||
doc_count=max(1, doc_count),
|
||||
selected_paths=paths_list or None,
|
||||
tiff_page_range=tiff_page_range,
|
||||
model=model or None,
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
|
||||
|
||||
@app.post("/batches", response_model=ExtractionBatchRead)
|
||||
@@ -375,6 +520,42 @@ def create_extraction_task(request: ExtractionRequest, background_tasks: Backgro
|
||||
detail="Envia fields. Usa POST /schemas para descubrir campos antes de crear la tarea.",
|
||||
)
|
||||
|
||||
if request.source == "desktop":
|
||||
documents = [doc.model_dump() for doc in request.documents or []]
|
||||
unique_radicados = {str(doc["radicado"]) for doc in documents}
|
||||
total = len(unique_radicados) if unique_radicados else (request.document_count or len(documents))
|
||||
title = f"Extraccion desktop ({total} documentos)"
|
||||
with session_scope() as session:
|
||||
batch = ExtractionBatch(
|
||||
status="running",
|
||||
root_path="desktop://upload",
|
||||
request=request.model_dump(),
|
||||
updated_at=utcnow(),
|
||||
)
|
||||
session.add(batch)
|
||||
session.flush()
|
||||
batch_id = batch.id
|
||||
task = Task(
|
||||
type="extraction",
|
||||
status="running",
|
||||
title=title,
|
||||
request=request.model_dump(),
|
||||
result={
|
||||
"source": "desktop",
|
||||
"extraction_batch_id": batch_id,
|
||||
"processed": 0,
|
||||
"total": total,
|
||||
"row_count": 0,
|
||||
"partials": {},
|
||||
"errors": [],
|
||||
},
|
||||
started_at=utcnow(),
|
||||
updated_at=utcnow(),
|
||||
)
|
||||
session.add(task)
|
||||
session.flush()
|
||||
return serialize_task(task)
|
||||
|
||||
with session_scope() as session:
|
||||
task = Task(
|
||||
type="extraction",
|
||||
@@ -393,6 +574,284 @@ def create_extraction_task(request: ExtractionRequest, background_tasks: Backgro
|
||||
return serialized
|
||||
|
||||
|
||||
async def _process_sync(
|
||||
task_id: str,
|
||||
batch_id: str,
|
||||
content: bytes,
|
||||
effective_name: str,
|
||||
radicado: str,
|
||||
doc_type: str | None,
|
||||
request: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Procesa un documento sincronicamente en un thread pool (modo sin Redis).
|
||||
|
||||
Estrategia de concurrencia:
|
||||
- La llamada a la LLM (lenta) se ejecuta sin ningún lock, permitiendo
|
||||
que todos los documentos llamen a OpenAI en paralelo.
|
||||
- Solo la escritura del resultado en la BD usa SELECT FOR UPDATE para
|
||||
evitar la race condition en el contador `processed`.
|
||||
"""
|
||||
import asyncio as _asyncio
|
||||
import logging as _logging
|
||||
import tempfile as _tempfile
|
||||
from domain.service import InvoiceProcessingService
|
||||
from domain.strategies import OpenAICompatibleExtractor
|
||||
from domain.batch_processor import BatchProcessor
|
||||
from domain.estimator import Estimator
|
||||
from api.desktop_ingest import _service_for_model, _paths_for_type
|
||||
from domain.profiles import get_profile
|
||||
from domain.prompts import PromptBuilder
|
||||
from domain.normalizer import normalize_to_schema
|
||||
|
||||
_log = _logging.getLogger("extraction.ingest")
|
||||
|
||||
ext = os.path.splitext(effective_name)[1] or ".pdf"
|
||||
with _tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
|
||||
tmp.write(content)
|
||||
temp_path = tmp.name
|
||||
|
||||
try:
|
||||
# ── Fase 1: llamada a la LLM (sin lock, en paralelo) ──────────────
|
||||
def _call_llm():
|
||||
model = request.get("model") or None
|
||||
svc = _service_for_model(processing_service, model)
|
||||
req = dict(request)
|
||||
fields = list(req.get("fields") or [])
|
||||
doc_type_arg = doc_type
|
||||
structure_mode = str(req.get("structure_mode") or "radicados")
|
||||
doc_types = list(req.get("doc_types") or [])
|
||||
prof = get_profile("invoice")
|
||||
selected_paths = _paths_for_type(
|
||||
fields,
|
||||
doc_type_arg if structure_mode == "radicados" and doc_types else None,
|
||||
)
|
||||
if doc_type_arg and structure_mode == "radicados" and doc_types:
|
||||
try:
|
||||
prompt = PromptBuilder.build_for_type(
|
||||
doc_type_arg, selected_paths, {}, profile_name=prof.get("name", "invoice")
|
||||
)
|
||||
except Exception:
|
||||
prompt = PromptBuilder.build(selected_paths, {}, profile_name=prof.get("name", "invoice"))
|
||||
else:
|
||||
instr = req.get("instruction")
|
||||
prompt = instr or PromptBuilder.build(selected_paths, {}, profile_name=prof.get("name", "invoice"))
|
||||
tiff_range = req.get("tiff_page_range")
|
||||
model_name = getattr(svc.extractor, "model", "?")
|
||||
_log.info("→ OpenAI [%s] radicado=%s filename=%s", model_name, radicado, effective_name)
|
||||
data, _usage = svc.extract_full_from_pdf(
|
||||
temp_path,
|
||||
instruction=prompt,
|
||||
tiff_page_range=tiff_range,
|
||||
)
|
||||
tok_in = (_usage or {}).get("prompt_tokens", "?")
|
||||
tok_out = (_usage or {}).get("completion_tokens", "?")
|
||||
_log.info("← OpenAI [%s] radicado=%s tokens=%s→%s", model_name, radicado, tok_in, tok_out)
|
||||
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
|
||||
if isinstance(norm, dict):
|
||||
normalized_obj = norm.get("normalized") or data
|
||||
return normalized_obj, selected_paths
|
||||
|
||||
normalized_obj, selected_paths = await _asyncio.to_thread(_call_llm)
|
||||
|
||||
# ── Fase 2: persistir resultado (con FOR UPDATE para serializar) ───
|
||||
def _persist():
|
||||
from infrastructure.serialization import flatten_extraction_result
|
||||
from infrastructure.database import ExtractionRow, ExtractionBatch, utcnow
|
||||
req = dict(request)
|
||||
fields = list(req.get("fields") or [])
|
||||
alias_map = dict(req.get("alias_map") or {})
|
||||
with session_scope() as session:
|
||||
task = session.get(Task, task_id, with_for_update=True)
|
||||
if task is None:
|
||||
return {"error": "task_not_found"}
|
||||
result = dict(task.result or {})
|
||||
fused = {"radicado": radicado, "_source_pdf": effective_name}
|
||||
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=effective_name,
|
||||
row_index=row_index,
|
||||
row_json=row,
|
||||
raw_json=fused,
|
||||
))
|
||||
result["row_count"] = row_index
|
||||
processed_now = int(result.get("processed") or 0) + 1
|
||||
result["processed"] = processed_now
|
||||
task.result = result
|
||||
task.updated_at = utcnow()
|
||||
session.flush()
|
||||
mark_task_completed_if_done(session, task, batch_id)
|
||||
session.flush() # persiste status=completed / finished_at si aplica
|
||||
_log.info(
|
||||
"✓ guardado radicado=%s processed=%d/%d status=%s",
|
||||
radicado, processed_now, int(result.get("total") or 0), task.status,
|
||||
)
|
||||
session.refresh(task)
|
||||
return {
|
||||
"radicado": radicado,
|
||||
"filename": effective_name,
|
||||
"processed": processed_now,
|
||||
"total": int(result.get("total") or 0),
|
||||
"task": serialize_task(task),
|
||||
}
|
||||
|
||||
return await _asyncio.to_thread(_persist)
|
||||
|
||||
finally:
|
||||
try:
|
||||
os.unlink(temp_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
async def _enqueue_async(
|
||||
task_id: str,
|
||||
batch_id: str,
|
||||
content: bytes,
|
||||
effective_name: str,
|
||||
radicado: str,
|
||||
doc_type: str | None,
|
||||
instruction: str | None,
|
||||
tiff_page_range: Any,
|
||||
) -> None:
|
||||
"""Guarda el archivo en el volumen compartido y encola el job en ARQ."""
|
||||
ext = os.path.splitext(effective_name)[1] or ".pdf"
|
||||
job_file = os.path.join(settings.queue_dir, f"{uuid.uuid4().hex}{ext}")
|
||||
with open(job_file, "wb") as fh:
|
||||
fh.write(content)
|
||||
pool = await get_arq_pool()
|
||||
await pool.enqueue_job(
|
||||
"process_document",
|
||||
task_id=task_id,
|
||||
batch_id=batch_id,
|
||||
queue_path=job_file,
|
||||
radicado=radicado,
|
||||
filename=effective_name,
|
||||
doc_type=doc_type,
|
||||
instruction=instruction,
|
||||
tiff_page_range=list(tiff_page_range) if tiff_page_range else None,
|
||||
)
|
||||
|
||||
|
||||
def _load_task_info(task_id: str) -> tuple[str, dict[str, Any], dict[str, Any]]:
|
||||
"""Carga batch_id, request y task serializado desde BD."""
|
||||
with session_scope() as session:
|
||||
task = session.get(Task, task_id)
|
||||
if task is None:
|
||||
raise HTTPException(status_code=404, detail="Tarea no encontrada")
|
||||
if task.status not in {"pending", "running"}:
|
||||
raise HTTPException(status_code=400, detail="La tarea ya no acepta documentos.")
|
||||
if str((task.request or {}).get("source") or "") != "desktop":
|
||||
raise HTTPException(status_code=400, detail="La tarea no fue creada en modo desktop.")
|
||||
batch_id = str((task.result or {}).get("extraction_batch_id") or "")
|
||||
if not batch_id:
|
||||
raise HTTPException(status_code=400, detail="La tarea no tiene lote asociado.")
|
||||
return batch_id, dict(task.request or {}), serialize_task(task)
|
||||
|
||||
|
||||
@app.post("/tasks/{task_id}/ingest")
|
||||
async def ingest_task_document(
|
||||
task_id: str,
|
||||
file: UploadFile = File(...),
|
||||
radicado: str = Form(...),
|
||||
filename: str | None = Form(None),
|
||||
doc_type: str | None = Form(None),
|
||||
) -> dict[str, Any]:
|
||||
"""Procesa o encola un documento.
|
||||
|
||||
- ASYNC_QUEUE=false (default): procesa sincronicamente, responde con la tarea actualizada.
|
||||
- ASYNC_QUEUE=true: encola en Redis/ARQ, responde de inmediato (sondea GET /tasks/{id}).
|
||||
"""
|
||||
batch_id, request, serialized = _load_task_info(task_id)
|
||||
effective_name = filename or file.filename or "documento.pdf"
|
||||
content = await file.read()
|
||||
|
||||
if settings.processing.async_queue:
|
||||
await _enqueue_async(
|
||||
task_id, batch_id, content, effective_name,
|
||||
radicado, doc_type, request.get("instruction"), request.get("tiff_page_range"),
|
||||
)
|
||||
return {"queued": True, "task": serialized}
|
||||
|
||||
# Modo sincrono: procesa ahora y devuelve resultado completo
|
||||
return await _process_sync(task_id, batch_id, content, effective_name, radicado, doc_type, request)
|
||||
|
||||
|
||||
@app.post("/tasks/{task_id}/ingest-many")
|
||||
async def ingest_task_documents_bulk(
|
||||
task_id: str,
|
||||
files: list[UploadFile] = File(...),
|
||||
radicados: str = Form(...),
|
||||
filenames: str | None = Form(None),
|
||||
doc_types: str | None = Form(None),
|
||||
) -> dict[str, Any]:
|
||||
"""Procesa o encola múltiples documentos en una sola petición HTTP.
|
||||
|
||||
Los parámetros `radicados`, `filenames` y `doc_types` son listas
|
||||
separadas por `|` que se corresponden posicionalmente con `files`.
|
||||
|
||||
En modo sincrono (ASYNC_QUEUE=false) procesa en paralelo con asyncio.
|
||||
En modo async (ASYNC_QUEUE=true) encola todos en ARQ.
|
||||
"""
|
||||
import asyncio as _asyncio
|
||||
|
||||
rad_list = radicados.split("|")
|
||||
fname_list = (filenames or "").split("|") if filenames else [""] * len(files)
|
||||
dtype_list = (doc_types or "").split("|") if doc_types else [None] * len(files)
|
||||
|
||||
if len(files) != len(rad_list):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Numero de archivos ({len(files)}) no coincide con radicados ({len(rad_list)})",
|
||||
)
|
||||
|
||||
batch_id, request, serialized = _load_task_info(task_id)
|
||||
contents = await _asyncio.gather(*[f.read() for f in files])
|
||||
|
||||
if settings.processing.async_queue:
|
||||
await _asyncio.gather(*[
|
||||
_enqueue_async(
|
||||
task_id, batch_id, contents[i],
|
||||
fname_list[i] or files[i].filename or f"doc_{i}.pdf",
|
||||
rad_list[i], dtype_list[i] or None,
|
||||
request.get("instruction"), request.get("tiff_page_range"),
|
||||
)
|
||||
for i in range(len(files))
|
||||
])
|
||||
return {"queued": len(files), "task": serialized}
|
||||
|
||||
# Modo sincrono: procesar todos en paralelo con semaforo para no saturar la API
|
||||
sem = _asyncio.Semaphore(settings.processing.concurrency_limit)
|
||||
|
||||
async def _bounded(i: int) -> dict[str, Any]:
|
||||
async with sem:
|
||||
return await _process_sync(
|
||||
task_id, batch_id, contents[i],
|
||||
fname_list[i] or files[i].filename or f"doc_{i}.pdf",
|
||||
rad_list[i], dtype_list[i] or None,
|
||||
request,
|
||||
)
|
||||
|
||||
results = await _asyncio.gather(*[_bounded(i) for i in range(len(files))], return_exceptions=True)
|
||||
errors = [str(r) for r in results if isinstance(r, Exception)]
|
||||
return {"processed": len(files) - len(errors), "errors": errors, "task": serialized}
|
||||
|
||||
|
||||
@app.get("/tasks", response_model=list[TaskRead])
|
||||
def list_tasks(type: str | None = None, status: str | None = None) -> list[dict[str, Any]]:
|
||||
with session_scope() as session:
|
||||
|
||||
@@ -89,7 +89,17 @@ class RetrySettings(BaseSettings):
|
||||
|
||||
class ProcessingSettings(BaseSettings):
|
||||
"""Configuración de procesamiento y concurrencia."""
|
||||
concurrency_limit: int = Field(100, validation_alias=AliasChoices("CONCURRENCY_LIMIT"))
|
||||
concurrency_limit: int = Field(50, validation_alias=AliasChoices("CONCURRENCY_LIMIT"))
|
||||
redis_url: str = Field("redis://localhost:6379/0", validation_alias=AliasChoices("REDIS_URL"))
|
||||
worker_concurrency: int = Field(25, validation_alias=AliasChoices("WORKER_CONCURRENCY"))
|
||||
# Límites OpenAI para el rate limiter distribuido
|
||||
openai_rpm_limit: int = Field(500, validation_alias=AliasChoices("OPENAI_RPM_LIMIT"))
|
||||
openai_tpm_limit: int = Field(800_000, validation_alias=AliasChoices("OPENAI_TPM_LIMIT"))
|
||||
# Directorio compartido entre API y workers para encolar archivos
|
||||
queue_dir: str = Field("/tmp/ips_queue", validation_alias=AliasChoices("QUEUE_DIR"))
|
||||
# False (default) = procesamiento sincrono sin Redis/workers.
|
||||
# True = cola ARQ distribuida (requiere Redis corriendo y workers activos).
|
||||
async_queue: bool = Field(False, validation_alias=AliasChoices("ASYNC_QUEUE"))
|
||||
model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8", extra="ignore")
|
||||
|
||||
|
||||
@@ -102,6 +112,23 @@ class AppSettings(BaseSettings):
|
||||
csv_delimiter: str = Field("|", validation_alias=AliasChoices("CSV_DELIMITER"))
|
||||
model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8", extra="ignore")
|
||||
|
||||
# Accesos directos para conveniencia
|
||||
@property
|
||||
def redis_url(self) -> str:
|
||||
return self.processing.redis_url
|
||||
|
||||
@property
|
||||
def openai_rpm_limit(self) -> int:
|
||||
return self.processing.openai_rpm_limit
|
||||
|
||||
@property
|
||||
def openai_tpm_limit(self) -> int:
|
||||
return self.processing.openai_tpm_limit
|
||||
|
||||
@property
|
||||
def queue_dir(self) -> str:
|
||||
return self.processing.queue_dir
|
||||
|
||||
|
||||
settings = AppSettings()
|
||||
|
||||
|
||||
@@ -110,6 +110,60 @@ class Estimator:
|
||||
def _accumulate_usage(self, agg, new_usage):
|
||||
accumulate_usage(agg, new_usage)
|
||||
|
||||
def estimate_from_file(
|
||||
self,
|
||||
file_path: str,
|
||||
doc_count: int = 1,
|
||||
selected_paths: Optional[List[str]] = None,
|
||||
tiff_page_range: Optional[tuple] = None,
|
||||
model: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Estima tokens y costos procesando un único archivo subido desde el cliente.
|
||||
|
||||
Si `model` se especifica y es distinto al modelo actual del extractor,
|
||||
crea un extractor temporal con ese modelo para obtener el uso real de tokens.
|
||||
"""
|
||||
from api.desktop_ingest import _service_for_model
|
||||
|
||||
service = _service_for_model(self.service, model)
|
||||
prof = get_profile('invoice')
|
||||
paths = selected_paths or []
|
||||
prompt = PromptBuilder.build(paths, {}, profile_name=prof.get('name', 'invoice'))
|
||||
|
||||
data, usage_ext = service.extract_full_from_pdf(
|
||||
file_path,
|
||||
instruction=prompt,
|
||||
tiff_page_range=tiff_page_range,
|
||||
)
|
||||
if isinstance(data, dict):
|
||||
data['radicado'] = 'sample'
|
||||
else:
|
||||
data = {"radicado": 'sample', "data": data}
|
||||
|
||||
norm = normalize_to_schema(data, paths=paths, profile_name=prof.get('name', 'invoice'))
|
||||
usage_norm = None
|
||||
if isinstance(norm, dict):
|
||||
usage_norm = norm.get("_usage")
|
||||
|
||||
result = self._build_result(usage_ext, usage_norm, doc_count, 'sample')
|
||||
|
||||
# Enriquecer con modelo efectivo y conteo de páginas para el frontend
|
||||
effective_model = model or getattr(service.extractor, 'model', None) or getattr(self.service.extractor, 'model', None)
|
||||
result['extractor_model'] = effective_model
|
||||
result['normalizer_model'] = settings.api.normalizer_model or 'gpt-5.4-mini'
|
||||
|
||||
# Estimar numero de paginas a partir de prompt_tokens de imagen:
|
||||
# GPT-4o cobra ~85 tok base + ~170 tok por tile 512x512.
|
||||
# Una A4 a 150 DPI (~1240x1754) genera ~6 tiles → ~1105 tokens.
|
||||
# Aproximamos: pages = max(1, round(image_tokens / 1000))
|
||||
prompt_tokens = int((usage_ext or {}).get('prompt_tokens') or 0)
|
||||
text_prompt_est = len(prompt) // 4 # 1 token ≈ 4 chars
|
||||
image_tokens_est = max(0, prompt_tokens - text_prompt_est)
|
||||
pages_est = max(1, round(image_tokens_est / 1000))
|
||||
result['pages_per_doc'] = pages_est
|
||||
|
||||
return result
|
||||
|
||||
def _build_result(self, usage_ext, usage_norm, total_docs, first_radicado):
|
||||
model_ext = getattr(self.service.extractor, "model", None)
|
||||
cost_ext = compute_cost_from_usage(usage_ext, "extract", model=model_ext)
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
"""Rate limiter distribuido para llamadas a OpenAI.
|
||||
|
||||
Usa Redis como backend para compartir el estado entre múltiples workers.
|
||||
Implementa el algoritmo de token bucket con ventana deslizante para respetar
|
||||
los límites de RPM (requests per minute) y TPM (tokens per minute) de OpenAI.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Límites por defecto para OpenAI gpt-4o (tier 1).
|
||||
# Sobreescribibles con variables de entorno vía config.
|
||||
DEFAULT_RPM = 500 # requests per minute
|
||||
DEFAULT_TPM = 800_000 # tokens per minute
|
||||
|
||||
|
||||
class DistributedRateLimiter:
|
||||
"""Token bucket distribuido respaldado por Redis.
|
||||
|
||||
Un solo objeto se comparte entre todos los workers de un proceso.
|
||||
Múltiples réplicas del worker comparten el estado en Redis.
|
||||
|
||||
Uso:
|
||||
limiter = DistributedRateLimiter(redis_url="redis://redis:6379/1")
|
||||
async with limiter.acquire(estimated_tokens=4000):
|
||||
response = await openai_client.chat.completions.create(...)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
redis_url: str,
|
||||
rpm: int = DEFAULT_RPM,
|
||||
tpm: int = DEFAULT_TPM,
|
||||
key_prefix: str = "ips:rl",
|
||||
) -> None:
|
||||
self._redis_url = redis_url
|
||||
self.rpm = rpm
|
||||
self.tpm = tpm
|
||||
self._prefix = key_prefix
|
||||
self._client: Optional[aioredis.Redis] = None
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def _get_client(self) -> aioredis.Redis:
|
||||
if self._client is None:
|
||||
async with self._lock:
|
||||
if self._client is None:
|
||||
self._client = aioredis.from_url(
|
||||
self._redis_url,
|
||||
encoding="utf-8",
|
||||
decode_responses=True,
|
||||
socket_connect_timeout=5,
|
||||
)
|
||||
return self._client
|
||||
|
||||
async def _sliding_window_count(self, key: str, window_s: float = 60.0) -> int:
|
||||
"""Cuenta eventos en la ventana deslizante usando un Sorted Set de Redis."""
|
||||
client = await self._get_client()
|
||||
now = time.monotonic()
|
||||
# Eliminar eventos expirados
|
||||
await client.zremrangebyscore(key, "-inf", now - window_s)
|
||||
return await client.zcard(key)
|
||||
|
||||
async def _record_event(self, key: str, tokens: int = 0, window_s: float = 60.0) -> None:
|
||||
"""Registra un evento en el sorted set. El score es el timestamp."""
|
||||
client = await self._get_client()
|
||||
now = time.monotonic()
|
||||
member = f"{now}:{tokens}"
|
||||
pipe = client.pipeline()
|
||||
pipe.zadd(key, {member: now})
|
||||
pipe.expire(key, int(window_s) + 5)
|
||||
await pipe.execute()
|
||||
|
||||
async def _token_sum_in_window(self, key: str, window_s: float = 60.0) -> int:
|
||||
"""Suma los tokens de todos los eventos en la ventana."""
|
||||
client = await self._get_client()
|
||||
now = time.monotonic()
|
||||
await client.zremrangebyscore(key, "-inf", now - window_s)
|
||||
members = await client.zrange(key, 0, -1)
|
||||
total = 0
|
||||
for m in members:
|
||||
parts = m.split(":", 1)
|
||||
if len(parts) == 2:
|
||||
try:
|
||||
total += int(parts[1])
|
||||
except ValueError:
|
||||
pass
|
||||
return total
|
||||
|
||||
async def acquire(self, estimated_tokens: int = 2000, max_wait_s: float = 120.0) -> None:
|
||||
"""Espera hasta que haya capacidad disponible en RPM y TPM.
|
||||
|
||||
Args:
|
||||
estimated_tokens: Tokens que se espera consumir (input + output).
|
||||
max_wait_s: Tiempo máximo de espera antes de lanzar TimeoutError.
|
||||
"""
|
||||
rpm_key = f"{self._prefix}:rpm"
|
||||
tpm_key = f"{self._prefix}:tpm"
|
||||
start = time.monotonic()
|
||||
sleep_s = 0.5
|
||||
|
||||
while True:
|
||||
rpm_count = await self._sliding_window_count(rpm_key)
|
||||
tpm_used = await self._token_sum_in_window(tpm_key)
|
||||
|
||||
rpm_ok = rpm_count < self.rpm
|
||||
tpm_ok = (tpm_used + estimated_tokens) <= self.tpm
|
||||
|
||||
if rpm_ok and tpm_ok:
|
||||
await self._record_event(rpm_key, tokens=0)
|
||||
await self._record_event(tpm_key, tokens=estimated_tokens)
|
||||
return
|
||||
|
||||
elapsed = time.monotonic() - start
|
||||
if elapsed >= max_wait_s:
|
||||
raise TimeoutError(
|
||||
f"Rate limiter: espera maxima alcanzada ({max_wait_s}s). "
|
||||
f"rpm={rpm_count}/{self.rpm} tpm={tpm_used}/{self.tpm}"
|
||||
)
|
||||
|
||||
reason = "RPM" if not rpm_ok else "TPM"
|
||||
logger.debug(
|
||||
"Rate limiter: esperando por %s (rpm=%d/%d tpm=%d/%d) — sleep %.1fs",
|
||||
reason, rpm_count, self.rpm, tpm_used, self.tpm, sleep_s,
|
||||
)
|
||||
await asyncio.sleep(sleep_s)
|
||||
sleep_s = min(sleep_s * 1.5, 10.0)
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._client is not None:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
|
||||
# Instancia global compartida por todos los workers del proceso.
|
||||
# Se inicializa de forma lazy en el primer uso.
|
||||
_limiter: Optional[DistributedRateLimiter] = None
|
||||
|
||||
|
||||
def get_rate_limiter() -> DistributedRateLimiter:
|
||||
"""Devuelve el rate limiter global, inicializándolo si es necesario."""
|
||||
global _limiter
|
||||
if _limiter is None:
|
||||
from domain.config import settings
|
||||
_limiter = DistributedRateLimiter(
|
||||
redis_url=settings.redis_url,
|
||||
rpm=settings.openai_rpm_limit,
|
||||
tpm=settings.openai_tpm_limit,
|
||||
)
|
||||
return _limiter
|
||||
@@ -0,0 +1,36 @@
|
||||
"""Helpers para subidas temporales desde la app desktop."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from contextlib import contextmanager
|
||||
from typing import Iterator
|
||||
|
||||
from fastapi import UploadFile
|
||||
|
||||
|
||||
@contextmanager
|
||||
def temp_upload_path(upload: UploadFile, *, suffix: str | None = None) -> Iterator[str]:
|
||||
"""Guarda un UploadFile en disco temporal y lo elimina al salir."""
|
||||
ext = suffix
|
||||
if ext is None and upload.filename:
|
||||
_, ext = os.path.splitext(upload.filename)
|
||||
if not ext:
|
||||
ext = ".bin"
|
||||
|
||||
fd, path = tempfile.mkstemp(suffix=ext)
|
||||
os.close(fd)
|
||||
try:
|
||||
upload.file.seek(0)
|
||||
with open(path, "wb") as handle:
|
||||
while True:
|
||||
chunk = upload.file.read(1024 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
handle.write(chunk)
|
||||
yield path
|
||||
finally:
|
||||
try:
|
||||
os.unlink(path)
|
||||
except OSError:
|
||||
pass
|
||||
@@ -1,5 +1,6 @@
|
||||
coloredlogs==15.0.1
|
||||
fastapi==0.115.6
|
||||
python-multipart==0.0.20
|
||||
openai==2.1.0
|
||||
pdf2image==1.17.0
|
||||
pillow==11.3.0
|
||||
@@ -10,3 +11,6 @@ python-dotenv==1.1.1
|
||||
SQLAlchemy==2.0.36
|
||||
tenacity==9.1.2
|
||||
uvicorn[standard]==0.32.1
|
||||
# Cola de tareas distribuida
|
||||
arq==0.26.3
|
||||
redis[hiredis]==5.2.1
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Worker ARQ para procesamiento paralelo de documentos."""
|
||||
@@ -0,0 +1,18 @@
|
||||
"""Punto de entrada del worker ARQ.
|
||||
|
||||
Ejecutar con:
|
||||
python -m worker.main
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from arq import run_worker
|
||||
from worker.settings import WorkerSettings
|
||||
from domain.logging_config import configure_logging
|
||||
|
||||
configure_logging()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if __name__ == "__main__":
|
||||
logger.info("Iniciando worker ARQ...")
|
||||
asyncio.run(run_worker(WorkerSettings))
|
||||
@@ -0,0 +1,84 @@
|
||||
"""Configuración del worker ARQ.
|
||||
|
||||
ARQ usa esta función `WorkerSettings` para saber qué funciones
|
||||
manejar, cuántos jobs en paralelo y cómo conectarse a Redis.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from arq.connections import RedisSettings
|
||||
|
||||
from domain.config import settings
|
||||
from domain.logging_config import configure_logging
|
||||
from domain.service import InvoiceProcessingService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def startup(ctx: dict[str, Any]) -> None:
|
||||
"""Inicializa el servicio de extracción compartido entre corrutinas del worker."""
|
||||
configure_logging()
|
||||
logger.info("Worker ARQ iniciando...")
|
||||
|
||||
from infrastructure.database import init_db
|
||||
from domain.config import extractor_llm_config
|
||||
from domain.strategies import OpenAICompatibleExtractor
|
||||
from domain.batch_processor import BatchProcessor
|
||||
from domain.estimator import Estimator
|
||||
init_db()
|
||||
|
||||
extractor_config = extractor_llm_config()
|
||||
extractor = OpenAICompatibleExtractor(
|
||||
api_key=str(extractor_config["api_key"]),
|
||||
base_url=str(extractor_config["base_url"]),
|
||||
model=str(extractor_config["model"]),
|
||||
)
|
||||
ctx["service"] = InvoiceProcessingService(
|
||||
extractor,
|
||||
batch_processor_factory=lambda s: BatchProcessor(s),
|
||||
estimator_factory=lambda s: Estimator(s),
|
||||
)
|
||||
logger.info("Worker ARQ listo. concurrency=%d", int(os.getenv("WORKER_CONCURRENCY", "25")))
|
||||
|
||||
|
||||
async def shutdown(ctx: dict[str, Any]) -> None:
|
||||
logger.info("Worker ARQ deteniendo...")
|
||||
from domain.rate_limiter import get_rate_limiter
|
||||
try:
|
||||
await get_rate_limiter().close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _parse_redis_settings() -> RedisSettings:
|
||||
url = settings.redis_url
|
||||
# Extraer host/port/db de la URL redis://host:port/db
|
||||
url = url.removeprefix("redis://")
|
||||
host_port, *db_parts = url.split("/")
|
||||
host, *port_parts = host_port.split(":")
|
||||
port = int(port_parts[0]) if port_parts else 6379
|
||||
db = int(db_parts[0]) if db_parts else 0
|
||||
return RedisSettings(host=host, port=port, database=db)
|
||||
|
||||
|
||||
class WorkerSettings:
|
||||
"""Configuración declarativa que ARQ espera en el módulo del worker."""
|
||||
from worker.tasks import process_document
|
||||
functions = [process_document]
|
||||
|
||||
redis_settings = _parse_redis_settings()
|
||||
|
||||
# Cuántos jobs corre en paralelo ESTE proceso worker
|
||||
max_jobs: int = int(os.getenv("WORKER_CONCURRENCY", "25"))
|
||||
|
||||
# Tiempo máximo por job antes de cancelarlo (segundos)
|
||||
job_timeout: int = int(os.getenv("EXTRACTOR_TIMEOUT_S", "120")) * 3
|
||||
|
||||
# Retener resultados en Redis 1 hora para diagnóstico
|
||||
keep_result: int = 3600
|
||||
|
||||
on_startup = startup
|
||||
on_shutdown = shutdown
|
||||
@@ -0,0 +1,236 @@
|
||||
"""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 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,
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
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].
|
||||
"""
|
||||
service = ctx["service"]
|
||||
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"}
|
||||
|
||||
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]:
|
||||
data, _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
|
||||
if isinstance(norm, dict):
|
||||
normalized_obj = norm.get("normalized") or data
|
||||
return raw_data, normalized_obj
|
||||
|
||||
_raw, normalized_obj = await asyncio.to_thread(_extract_and_normalize)
|
||||
|
||||
# Persistir resultado en BD
|
||||
def _persist() -> dict[str, Any]:
|
||||
with session_scope() as session:
|
||||
task = session.get(Task, task_id)
|
||||
if task is None:
|
||||
return {"skipped": True, "reason": "task_not_found"}
|
||||
|
||||
result = dict(task.result or {})
|
||||
|
||||
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
|
||||
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)
|
||||
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
|
||||
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)
|
||||
Reference in New Issue
Block a user