Spaces:
Sleeping
Sleeping
| """ | |
| TradeFlow AI — Batch Processing Tasks (T-072) | |
| """ | |
| from __future__ import annotations | |
| import contextlib | |
| import logging | |
| logger = logging.getLogger("tasks.batch") | |
| def process_batch(batch_id: str) -> dict: | |
| """ | |
| Main Celery task: invoke the LangGraph pipeline for a batch. | |
| Called by POST /api/v1/batches after file upload. | |
| """ | |
| try: | |
| import asyncio | |
| from packages.agents.src.graph import get_graph # type: ignore | |
| from ..config import settings # type: ignore | |
| graph = get_graph() | |
| config = {"configurable": {"thread_id": batch_id}} | |
| # Fetch batch documents from DB | |
| state = _load_initial_state(batch_id, settings) | |
| # Run the graph synchronously from Celery worker | |
| result = asyncio.run(graph.ainvoke(state, config=config)) | |
| logger.info(f"Batch {batch_id} pipeline complete. Status: {result.get('ceisa_response', {}).get('status')}") | |
| return {"batch_id": batch_id, "status": "complete"} | |
| except Exception as e: | |
| logger.error(f"Batch {batch_id} pipeline failed: {e}") | |
| _mark_batch_failed(batch_id, str(e)) | |
| raise | |
| def cleanup_expired_batches() -> int: | |
| """Delete batches older than 7 days (Celery beat hourly).""" | |
| import asyncio | |
| try: | |
| return asyncio.run(_cleanup_async()) | |
| except Exception as e: | |
| logger.error(f"Batch cleanup failed: {e}") | |
| return 0 | |
| async def _cleanup_async() -> int: | |
| from ..db.database import get_async_session # type: ignore | |
| async with get_async_session() as db: | |
| result = await db.execute( | |
| "DELETE FROM batches WHERE created_at < NOW() - INTERVAL '7 days' " | |
| "AND status NOT IN ('ACCEPTED', 'REJECTED') RETURNING id" | |
| ) | |
| deleted = result.rowcount if hasattr(result, "rowcount") else 0 | |
| logger.info(f"Cleaned up {deleted} expired batches") | |
| return deleted | |
| def _load_initial_state(batch_id: str, settings) -> dict: | |
| """Load initial DeclarationState from database.""" | |
| import asyncio | |
| async def _fetch(): | |
| from ..db.database import get_async_session # type: ignore | |
| async with get_async_session() as db: | |
| rows = await db.execute( | |
| "SELECT id, doc_type, storage_path, original_name FROM documents WHERE batch_id = $1", | |
| batch_id, | |
| ) | |
| docs = [dict(r) for r in rows] | |
| tier_row = await db.execute( | |
| "SELECT tier FROM batches WHERE id = $1", batch_id | |
| ) | |
| tier = "sme" | |
| if tier_row: | |
| tier = tier_row[0].get("tier", "sme") | |
| return { | |
| "batch_id": batch_id, | |
| "tier": tier, | |
| "documents": docs, | |
| "messages": [], | |
| "operator_corrections": [], | |
| "submission_attempt": 0, | |
| "error": None, | |
| } | |
| return asyncio.run(_fetch()) | |
| def _mark_batch_failed(batch_id: str, error: str) -> None: | |
| import asyncio | |
| async def _update(): | |
| from ..db.database import get_async_session # type: ignore | |
| async with get_async_session() as db: | |
| await db.execute( | |
| "UPDATE batches SET status = 'FAILED', error_message = $1 WHERE id = $2", | |
| error[:500], batch_id, | |
| ) | |
| with contextlib.suppress(Exception): | |
| asyncio.run(_update()) | |