"""Servidor de tracking de IPSExtractor. Persiste el historial de lotes en SQLite y distribuye actualizaciones en tiempo real vía WebSockets al dashboard web. """ from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.encoders import jsonable_encoder from fastapi.responses import FileResponse from contextlib import asynccontextmanager from typing import Dict, List, Any, Optional from pydantic import BaseModel import uuid import os from datetime import datetime from .schemas.batch import BatchCreate, BatchUpdate, BatchRead, BatchStatus from .database import init_db, insert_batch, update_batch, get_all_batches, get_batch, insert_log, get_logs @asynccontextmanager async def lifespan(app: FastAPI): init_db() yield app = FastAPI(title="IPSExtractor Tracking API", lifespan=lifespan) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Cache en memoria para lotes activos (para broadcast rápido) active_batches: Dict[str, Dict[str, Any]] = {} active_connections: List[WebSocket] = [] async def broadcast(message: dict): json_compatible_message = jsonable_encoder(message) disconnected = [] for connection in active_connections: try: await connection.send_json(json_compatible_message) except Exception: disconnected.append(connection) for conn in disconnected: active_connections.remove(conn) @app.post("/batches", response_model=BatchRead) async def create_batch(batch: BatchCreate): batch_id = str(uuid.uuid4()) now = datetime.now() new_batch = { "id": batch_id, **batch.model_dump(), "status": BatchStatus.PROCESSING, "current_radicado_index": 0, "success_count": 0, "error_count": 0, "total_tokens": 0, "total_cost_usd": 0.0, "last_error": None, "created_at": now, "updated_at": now, } active_batches[batch_id] = new_batch insert_batch(new_batch) await broadcast({"event": "batch_created", "data": new_batch}) return new_batch @app.get("/batches", response_model=List[BatchRead]) async def list_batches(): # Retorna historial completo desde SQLite db_batches = get_all_batches() # Sobreescribe con datos en memoria si el lote está activo (más fresco) result = [] for b in db_batches: if b["id"] in active_batches: result.append(active_batches[b["id"]]) else: result.append(b) return result @app.get("/batches/{batch_id}", response_model=BatchRead) async def get_batch_endpoint(batch_id: str): if batch_id in active_batches: return active_batches[batch_id] batch = get_batch(batch_id) if not batch: raise HTTPException(status_code=404, detail="Lote no encontrado") return batch @app.patch("/batches/{batch_id}", response_model=BatchRead) async def update_batch_endpoint(batch_id: str, update: BatchUpdate): if batch_id in active_batches: batch = active_batches[batch_id] update_data = update.model_dump(exclude_unset=True) for key, value in update_data.items(): batch[key] = value batch["updated_at"] = datetime.now() # Si el lote se completó, sacarlo del cache activo if batch.get("status") in (BatchStatus.COMPLETED, BatchStatus.FAILED): active_batches.pop(batch_id, None) update_batch(batch_id, update_data) await broadcast({"event": "batch_updated", "data": batch}) return batch # Lote no está en memoria (completado), actualizar solo en DB batch = get_batch(batch_id) if not batch: raise HTTPException(status_code=404, detail="Lote no encontrado") update_data = update.model_dump(exclude_unset=True) update_batch(batch_id, update_data) batch.update(update_data) return batch class LogEntry(BaseModel): message: str level: str = "info" radicado: Optional[str] = None class LogRead(BaseModel): id: int batch_id: str timestamp: str level: str radicado: Optional[str] message: str @app.post("/batches/{batch_id}/logs", status_code=201) async def add_log(batch_id: str, entry: LogEntry): """Añade una entrada de log para un lote.""" insert_log(batch_id, entry.message, entry.level, entry.radicado) # Broadcast log también por WebSocket para actualizaciones en vivo await broadcast({"event": "log_added", "batch_id": batch_id, "data": { "level": entry.level, "radicado": entry.radicado, "message": entry.message }}) return {"ok": True} @app.get("/batches/{batch_id}/logs", response_model=List[LogRead]) async def list_logs(batch_id: str, limit: int = 200): """Retorna el historial de logs de un lote.""" return get_logs(batch_id, limit) @app.get("/download/{batch_id}") async def download_csv(batch_id: str): """Sirve el archivo CSV generado por el lote de extracción.""" batch = get_batch(batch_id) if not batch: raise HTTPException(status_code=404, detail="Lote no encontrado") csv_path = batch.get("csv_output_path") if not csv_path: raise HTTPException(status_code=404, detail="Este lote no tiene un archivo CSV asociado") if not os.path.isfile(csv_path): raise HTTPException(status_code=404, detail=f"Archivo CSV no encontrado en el servidor: {csv_path}") filename = os.path.basename(csv_path) return FileResponse( path=csv_path, media_type="text/csv", filename=filename, ) @app.websocket("/ws/tracking") async def websocket_endpoint(websocket: WebSocket): await websocket.accept() active_connections.append(websocket) try: while True: await websocket.receive_text() except WebSocketDisconnect: if websocket in active_connections: active_connections.remove(websocket)