Spaces:
Sleeping
Sleeping
File size: 3,381 Bytes
db4ba8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | """
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())
|