Files
IPSExtractor/services/extraction/api/main.py
T

328 lines
9.9 KiB
Python

"""API REST del microservicio de extraccion."""
from __future__ import annotations
from contextlib import asynccontextmanager
from typing import Any
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from sqlalchemy import select
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
from domain.estimator import Estimator
from domain.service import InvoiceProcessingService
from domain.strategies import OpenAICompatibleExtractor
from infrastructure.database import (
ExtractionBatch,
ExtractionRow,
init_db,
session_scope,
utcnow,
)
from infrastructure.serialization import flatten_extraction_result
def build_processing_service() -> InvoiceProcessingService:
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"]),
)
return InvoiceProcessingService(
extractor,
batch_processor_factory=lambda service: BatchProcessor(service),
estimator_factory=lambda service: Estimator(service),
)
processing_service = build_processing_service()
@asynccontextmanager
async def lifespan(app: FastAPI):
init_db()
yield
app = FastAPI(
title="IPSExtractor Extraction Service",
version="1.0.0",
lifespan=lifespan,
)
class DocumentType(BaseModel):
name: str
keyword: str = ""
class SchemaRequest(BaseModel):
document_path: str | None = None
root_path: str | None = None
doc_types: list[DocumentType] | None = None
instruction: str | None = None
tiff_page_range: tuple[int, int] | None = None
structure_mode: str = "radicados"
class ExtractionRequest(BaseModel):
root_path: str
keywords: list[str] | None = None
doc_types: list[DocumentType] | None = None
fields: list[str] = Field(default_factory=list)
alias_map: dict[str, str] = Field(default_factory=dict)
instruction: str | None = None
tiff_page_range: tuple[int, int] | None = None
structure_mode: str = "radicados"
retry_on_empty: bool = False
radicado_filter: list[str] | None = None
class EstimateRequest(BaseModel):
root_path: str
keywords: list[str] | None = None
doc_types: list[DocumentType] | None = None
schema_map: dict[str, str] | None = Field(default=None, alias="schema")
selected_fields: list[str] | None = None
schema_by_type: dict[str, dict[str, str]] | None = None
selected_paths: list[str] | None = None
structure_mode: str = "radicados"
class ExtractionBatchRead(BaseModel):
id: str
status: str
root_path: str | None
request: dict[str, Any]
error: str | None
created_at: str
updated_at: str
def serialize_batch(batch: ExtractionBatch) -> dict[str, Any]:
return {
"id": batch.id,
"status": batch.status,
"root_path": batch.root_path,
"request": batch.request,
"error": batch.error,
"created_at": batch.created_at.isoformat(),
"updated_at": batch.updated_at.isoformat(),
}
@app.get("/health")
def health() -> dict[str, str]:
return {"status": "ok"}
@app.get("/config")
def get_config() -> dict[str, Any]:
extractor_config = extractor_llm_config()
normalizer_config = normalizer_llm_config()
return {
"extractor": {
"provider": extractor_config["provider"],
"base_url": extractor_config["base_url"],
"model": extractor_config["model"],
"timeout_s": extractor_config["timeout_s"],
"api_key_configured": bool(str(extractor_config["api_key"]).strip()),
},
"normalizer": {
"provider": normalizer_config["provider"],
"base_url": normalizer_config["base_url"],
"model": normalizer_config["model"],
"timeout_s": normalizer_config["timeout_s"],
"api_key_configured": bool(str(normalizer_config["api_key"]).strip()),
},
"processing": {
"concurrency_limit": settings.processing.concurrency_limit,
"allowed_document_types": list(get_allowed_document_types()),
},
"imaging": {
"pdf_dpi": settings.imaging.pdf_dpi,
"pdf_thread_count": settings.imaging.pdf_thread_count,
"tiff_pages": settings.imaging.tiff_pages,
"poppler_path": settings.imaging.poppler_path,
},
}
@app.get("/models")
def list_models() -> dict[str, Any]:
extractor_config = extractor_llm_config()
normalizer_config = normalizer_llm_config()
try:
models = processing_service.extractor.llm.list_models()
return {
"available": models,
"configured_extractor_model": extractor_config["model"],
"configured_normalizer_model": normalizer_config["model"],
"source": extractor_config["base_url"],
}
except Exception as exc:
return {
"available": [],
"configured_extractor_model": extractor_config["model"],
"configured_normalizer_model": normalizer_config["model"],
"source": extractor_config["base_url"],
"warning": f"No se pudieron consultar modelos del backend: {exc}",
}
@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],
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,
}
raise HTTPException(
status_code=400,
detail="Envia document_path o root_path con doc_types para descubrir campos.",
)
@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,
)
@app.post("/batches", response_model=ExtractionBatchRead)
def run_extraction(request: ExtractionRequest) -> dict[str, Any]:
if not request.fields:
raise HTTPException(
status_code=400,
detail="El API no tiene selector visual: envia fields. Usa POST /schemas para descubrirlos.",
)
with session_scope() as session:
batch = ExtractionBatch(
status="running",
root_path=request.root_path,
request=request.model_dump(),
updated_at=utcnow(),
)
session.add(batch)
session.flush()
batch_id = batch.id
try:
results = processing_service.process_root_with_field_selection(
request.root_path,
request.keywords,
doc_types=[dt.model_dump() for dt in request.doc_types or []] or None,
fields=request.fields,
csv_output_path=None,
instruction_override=request.instruction,
alias_map=request.alias_map,
tiff_page_range=request.tiff_page_range,
structure_mode=request.structure_mode,
retry_on_empty=request.retry_on_empty,
radicado_filter=request.radicado_filter,
)
with session_scope() as session:
row_index = 0
for item in results:
raw = item if isinstance(item, dict) else {"data": item}
rows = flatten_extraction_result(raw, request.fields, request.alias_map)
for row in rows:
row_index += 1
session.add(
ExtractionRow(
batch_id=batch_id,
radicado=str(raw.get("radicado") or row.get("radicado") or ""),
source_pdf=raw.get("_source_pdf"),
row_index=row_index,
row_json=row,
raw_json=raw,
)
)
batch = session.get(ExtractionBatch, batch_id)
if batch is None:
raise RuntimeError("No se encontro el lote recien creado.")
batch.status = "completed"
batch.updated_at = utcnow()
session.flush()
return serialize_batch(batch)
except Exception as exc:
with session_scope() as session:
batch = session.get(ExtractionBatch, batch_id)
if batch is not None:
batch.status = "failed"
batch.error = str(exc)
batch.updated_at = utcnow()
raise HTTPException(status_code=500, detail=str(exc)) from exc
@app.get("/batches")
def list_batches() -> list[dict[str, Any]]:
with session_scope() as session:
batches = session.scalars(
select(ExtractionBatch).order_by(ExtractionBatch.created_at.desc())
).all()
return [serialize_batch(batch) for batch in batches]
@app.get("/batches/{batch_id}")
def get_batch(batch_id: str) -> dict[str, Any]:
with session_scope() as session:
batch = session.get(ExtractionBatch, batch_id)
if batch is None:
raise HTTPException(status_code=404, detail="Lote no encontrado")
return serialize_batch(batch)
@app.get("/batches/{batch_id}/rows")
def get_rows(batch_id: str, limit: int = 500, offset: int = 0) -> dict[str, Any]:
with session_scope() as session:
batch = session.get(ExtractionBatch, batch_id)
if batch is None:
raise HTTPException(status_code=404, detail="Lote no encontrado")
rows = session.scalars(
select(ExtractionRow)
.where(ExtractionRow.batch_id == batch_id)
.order_by(ExtractionRow.row_index, ExtractionRow.id)
.offset(offset)
.limit(limit)
).all()
return {
"batch_id": batch_id,
"rows": [row.row_json for row in rows],
"limit": limit,
"offset": offset,
}