"""商品 AI 异步任务（Celery）：批量补全。 状态与进度写入 AsyncTaskLog，供 GET /api/admin/ai-tasks/{id} 轮询。""" import asyncio from datetime import datetime, timezone from contextlib import asynccontextmanager from sqlalchemy import select, func, update from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession from . import ai_enrich as ae from app.celery_app import celery_app from app.config import settings from app.database import AsyncSessionLocal from app.core.models.task import AsyncTaskLog, AiTaskItem @asynccontextmanager async def _task_engine(): """Celery 任务专用：在本次 asyncio.run 的事件循环内新建独立引擎，用完即弃。 避免复用绑定在 uvicorn 循环上的全局引擎导致 'Future attached to a different loop'。""" engine = create_async_engine(settings.DATABASE_URL, pool_pre_ping=True) try: yield async_sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False, autoflush=False) finally: await engine.dispose() def _dt(value): return value.isoformat() if value else None def _serialize_task_summary(row) -> dict: result = row.result or {} payload = row.payload or {} return { "task_id": row.id, "task_type": row.task_type, "status": row.status, "payload": payload, "progress": result.get("progress") or {"done": 0, "total": payload.get("count", 0)}, "error": result.get("error"), "created_at": _dt(getattr(row, "created_at", None)), "completed_at": _dt(getattr(row, "completed_at", None)), } def _serialize_task_row(row) -> dict: return { "task_id": row.id, "task_type": row.task_type, "status": row.status, "result": row.result, } async def _create_task_items(task_id: int, tenant_id: int, rows: list[dict]) -> None: async with AsyncSessionLocal() as db: db.add_all([ AiTaskItem( task_id=task_id, tenant_id=tenant_id, product_id=row.get("id") or row.get("product_id"), row_index=i, status="pending", source_row=row, ) for i, row in enumerate(rows) ]) await db.commit() async def _create_task_row(tenant_id: int, task_type: str, payload: dict) -> int: """同步创建一条 pending 任务，返回 id 给前端；随后再 .delay() 触发执行。""" async with AsyncSessionLocal() as db: row = AsyncTaskLog( tenant_id=tenant_id, task_type=task_type, status="pending", payload=payload, ) db.add(row) await db.commit() await db.refresh(row) return row.id async def _update_task(task_id: int, *, status: str | None = None, result: dict | None = None, session_factory=None) -> None: async with (session_factory or AsyncSessionLocal)() as db: row = await db.get(AsyncTaskLog, task_id) if row is None: return if status is not None: row.status = status if status in ("success", "failed", "cancelled"): row.completed_at = datetime.now(timezone.utc) if result is not None: row.result = result await db.commit() async def _mark_failed(task_id: int, exc: Exception) -> None: """失败上报：也用独立引擎，避免在新循环里碰全局引擎。""" async with _task_engine() as SF: await _update_task(task_id, status="failed", result={"error": str(exc)[:300]}, session_factory=SF) async def _is_cancelled(task_id: int, session_factory=None) -> bool: async with (session_factory or AsyncSessionLocal)() as db: row = await db.get(AsyncTaskLog, task_id) return row is not None and row.status == "cancelled" async def _enrich_rows_collect(rows: list, extra: dict, progress_cb, should_continue=None) -> list: """逐行调用 AI 补全，每行完成后上报进度。 单行 AI 失败（如 529 过载）只保留原行、继续，不拖垮整批。""" out = [] for i, row in enumerate(rows): if should_continue is not None: keep_going = should_continue() if asyncio.iscoroutine(keep_going): keep_going = await keep_going if not keep_going: break try: out.append(await ae.enrich_one_pure(row, extra, search_context="", store_context="")) except Exception: out.append(row) # 保留原行，该行不补全 if progress_cb is not None: await progress_cb(i + 1, len(rows)) return out def _task_item_select(task_id: int, tenant_id: int, *statuses: str): query = select(AiTaskItem).where( AiTaskItem.task_id == task_id, AiTaskItem.tenant_id == tenant_id, ) if statuses: query = query.where(AiTaskItem.status.in_(statuses)) return query.order_by(AiTaskItem.row_index) def _retry_items_stmt(task_id: int, tenant_id: int): return ( update(AiTaskItem) .where( AiTaskItem.task_id == task_id, AiTaskItem.tenant_id == tenant_id, AiTaskItem.status == "failed", ) .values(status="pending", error=None) ) def _rows_from_items(items) -> list[dict]: return [ item.result_row if item.result_row is not None and item.status in ("success", "failed") else item.source_row for item in sorted(items, key=lambda item: item.row_index) ] @celery_app.task(bind=True, name="ai.enrich_batch", acks_late=True) def enrich_batch_task(self, *, task_id: int, tenant_id: int): async def run(): async with _task_engine() as SF: if await _is_cancelled(task_id, SF): return async with SF() as db: extra = await ae.get_ai_extra(db, tenant_id) total = (await db.execute(select(func.count()).select_from(AiTaskItem).where(AiTaskItem.task_id == task_id, AiTaskItem.tenant_id == tenant_id))).scalar_one() await _update_task(task_id, status="running", result={"progress": {"done": 0, "total": total}}, session_factory=SF) done = 0 failed = 0 # ponytail: 只处理 pending 的行，跳过已完成的（支持断点续跑） async with SF() as db: result = await db.execute( _task_item_select(task_id, tenant_id, "pending") ) items = result.scalars().all() # 已完成的行数（断点续跑时 > 0） done = total - len(items) CONCURRENCY = 20 # ponytail: 并发数，被限流再降 async def _do_one(item): nonlocal done, failed try: enriched = await ae.enrich_one_pure(item.source_row, extra, search_context="", store_context="") item_status = "success" error = None except Exception as exc: enriched = item.source_row item_status = "failed" error = str(exc)[:500] failed += 1 done += 1 async with SF() as db: current = await db.get(AiTaskItem, item.id) if current is not None: current.status = item_status current.result_row = enriched current.error = error await db.commit() # 按 CONCURRENCY 分批并发 for i in range(0, len(items), CONCURRENCY): if await _is_cancelled(task_id, SF): return batch = items[i:i + CONCURRENCY] await asyncio.gather(*[_do_one(item) for item in batch]) await _update_task( task_id, result={"progress": {"done": done, "total": total}, "failed": failed, "item_storage": "ai_task_items"}, session_factory=SF, ) if await _is_cancelled(task_id, SF): return await _update_task( task_id, status="success", result={"progress": {"done": done, "total": total}, "failed": failed, "item_storage": "ai_task_items"}, session_factory=SF, ) try: asyncio.run(run()) except Exception as exc: asyncio.run(_mark_failed(task_id, exc)) raise def _initial_commit_totals(result: dict | None, *, total: int, pending: int): result = result or {} return { "created": result.get("created", 0), "updated": result.get("updated", 0), "skipped": result.get("skipped", 0), "errors": [], }, total - pending @celery_app.task(bind=True, name="bulk_import.commit", acks_late=True) def commit_batch_task(self, *, task_id: int, tenant_id: int, user_id: int, mode: str = "skip", create_missing_taxonomy: bool = False): async def run(): from app.core.models.user import User from app.plugins.bulk_import import services as bi async with _task_engine() as SF: async with SF() as db: user = await db.get(User, user_id) if user is None or user.tenant_id != tenant_id: raise ValueError("import user not found") filters = (AiTaskItem.task_id == task_id, AiTaskItem.tenant_id == tenant_id) total = (await db.execute(select(func.count()).select_from(AiTaskItem).where(*filters))).scalar_one() pending = (await db.execute(select(func.count()).select_from(AiTaskItem).where(*filters, AiTaskItem.status == "pending"))).scalar_one() parent = await db.get(AsyncTaskLog, task_id) pending_items = (await db.execute(_task_item_select(task_id, tenant_id, "pending"))).scalars().all() prepared_rows = await bi.prepare_taxonomy( db, tenant_id, [dict(item.source_row) for item in pending_items], create_missing_taxonomy ) for item, prepared in zip(pending_items, prepared_rows): item.source_row = prepared await db.commit() from app.core.models.product import Product sku_rows = await db.execute(select(Product.sku, Product.id).where(Product.tenant_id == tenant_id)) existing = {sku: product_id for sku, product_id in sku_rows.all() if sku} brand_cache = await bi._build_brand_cache(db, tenant_id) totals, done = _initial_commit_totals(parent.result if parent else None, total=total, pending=pending) await _update_task(task_id, status="running", result={**totals, "progress": {"done": done, "total": total}}, session_factory=SF) while True: async with SF() as db: items = (await db.execute(_task_item_select(task_id, tenant_id, "pending").limit(100))).scalars().all() if not items or await _is_cancelled(task_id, SF): break for item in items: result = await bi.commit_rows( SF, user, [item.source_row], mode=mode, existing=existing, brand_cache=brand_cache, ) error = (result.get("errors") or [None])[0] status = "failed" if error else ("skipped" if result.get("skipped") else "success") async with SF() as db: current = (await db.execute( _task_item_select(task_id, tenant_id).where(AiTaskItem.id == item.id) )).scalar_one_or_none() if current is not None: current.status = status current.result_row = item.source_row current.error = (error or {}).get("error") if error else None await db.commit() for key in ("created", "updated", "skipped"): totals[key] += result.get(key, 0) if error: totals["errors"].append({"index": item.row_index, "error": error.get("error", "")[:200]}) done += 1 # 递增计数：包含断点续跑前已处理的行（含历史失败行） await _update_task(task_id, result={**totals, "progress": {"done": done, "total": total}}, session_factory=SF) if not await _is_cancelled(task_id, SF): await _update_task(task_id, status="success", result={**totals, "progress": {"done": done, "total": total}}, session_factory=SF) try: asyncio.run(run()) except Exception as exc: asyncio.run(_mark_failed(task_id, exc)) raise