| from __future__ import annotations |
| import asyncio, json, uuid, time |
| from sqlalchemy import select |
| from app.models.batch_job import BatchJob |
| from app.services.storage import async_session_maker |
| from app.services.connection_pool import get_client |
|
|
| MAX_BATCH_SIZE = 50 |
| MAX_CONCURRENT = 5 |
|
|
|
|
| async def create_batch(session, requests: list, model: str) -> dict: |
| if len(requests) > MAX_BATCH_SIZE: |
| requests = requests[:MAX_BATCH_SIZE] |
|
|
| batch_id = str(uuid.uuid4()) |
| total = len(requests) |
| results = [] |
| completed = 0 |
| failed = 0 |
|
|
| session.add(BatchJob( |
| id=batch_id, status="processing", total=total, |
| completed=0, failed=0, results="[]", |
| )) |
| await session.commit() |
|
|
| semaphore = asyncio.Semaphore(MAX_CONCURRENT) |
| client = get_client() |
|
|
| async def _process(req: dict, idx: int): |
| nonlocal completed, failed |
| async with semaphore: |
| try: |
| url = req.get("url", "") |
| method = req.get("method", "POST").upper() |
| headers = req.get("headers", {}) |
| body = req.get("body", {}) |
|
|
| if method == "GET": |
| resp = await client.get(url, headers=headers, timeout=120.0) |
| else: |
| resp = await client.post(url, headers=headers, json=body, timeout=120.0) |
|
|
| result = { |
| "index": idx, "status_code": resp.status_code, |
| "body": resp.text[:10000], |
| "error": "", |
| } |
| results.append(result) |
| if resp.is_success: |
| completed += 1 |
| else: |
| failed += 1 |
| except Exception as e: |
| results.append({"index": idx, "status_code": 0, "body": "", "error": str(e)}) |
| failed += 1 |
|
|
| tasks = [_process(req, i) for i, req in enumerate(requests)] |
| await asyncio.gather(*tasks) |
|
|
| status = "completed" if failed == 0 else ("partial" if completed > 0 else "failed") |
| await session.execute( |
| select(BatchJob).where(BatchJob.id == batch_id) |
| ) |
| row = (await session.execute(select(BatchJob).where(BatchJob.id == batch_id))).scalar_one() |
| row.status = status |
| row.completed = completed |
| row.failed = failed |
| row.results = json.dumps(results) |
| row.completed_at = time.time() |
| await session.commit() |
|
|
| return { |
| "batch_id": batch_id, |
| "status": status, |
| "total": total, |
| "completed": completed, |
| "failed": failed, |
| } |
|
|
|
|
| async def get_batch_history(session) -> list: |
| result = await session.execute( |
| select(BatchJob).order_by(BatchJob.created_at.desc()).limit(50) |
| ) |
| rows = result.scalars().all() |
| return [ |
| { |
| "batch_id": r.id, "status": r.status, |
| "total": r.total, "completed": r.completed, "failed": r.failed, |
| "created_at": r.created_at, "completed_at": r.completed_at, |
| } |
| for r in rows |
| ] |
|
|
