""" 项目向量化后台任务服务 """ import logging import uuid from typing import Optional from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql import func from app.core.database import AsyncSessionLocal from app.models.project import Project from app.models.project_vectorization_task import ProjectVectorizationTask from app.services.zvec_service import zvec_service logger = logging.getLogger(__name__) class ProjectVectorizationTaskService: """管理项目文件夹向量化任务的创建、执行和查询。""" RUNNING_STATUSES = ("pending", "running") @staticmethod def serialize_task(task: ProjectVectorizationTask) -> dict: return { "task_id": task.task_id, "project_id": task.project_id, "task_type": task.task_type, "status": task.status, "total": task.total, "processed": task.processed, "skipped": task.skipped, "failed": task.failed, "error_message": task.error_message, "created_at": task.created_at.isoformat() if task.created_at else None, "started_at": task.started_at.isoformat() if task.started_at else None, "finished_at": task.finished_at.isoformat() if task.finished_at else None, "updated_at": task.updated_at.isoformat() if task.updated_at else None, } async def create_task( self, db: AsyncSession, project_id: int, user_id: int, *, force: bool, ) -> ProjectVectorizationTask: task = ProjectVectorizationTask( task_id=uuid.uuid4().hex, project_id=project_id, user_id=user_id, task_type="full" if force else "incremental", status="pending", ) db.add(task) await db.commit() await db.refresh(task) return task async def get_running_task( self, db: AsyncSession, project_id: int, ) -> Optional[ProjectVectorizationTask]: result = await db.execute( select(ProjectVectorizationTask) .where( ProjectVectorizationTask.project_id == project_id, ProjectVectorizationTask.status.in_(self.RUNNING_STATUSES), ) .order_by(ProjectVectorizationTask.created_at.desc()) .limit(1) ) return result.scalar_one_or_none() async def get_task( self, db: AsyncSession, project_id: int, task_id: str, ) -> Optional[ProjectVectorizationTask]: result = await db.execute( select(ProjectVectorizationTask).where( ProjectVectorizationTask.project_id == project_id, ProjectVectorizationTask.task_id == task_id, ) ) return result.scalar_one_or_none() async def get_latest_task( self, db: AsyncSession, project_id: int, ) -> Optional[ProjectVectorizationTask]: result = await db.execute( select(ProjectVectorizationTask) .where(ProjectVectorizationTask.project_id == project_id) .order_by(ProjectVectorizationTask.created_at.desc()) .limit(1) ) return result.scalar_one_or_none() async def run_task(self, task_id: str) -> None: async with AsyncSessionLocal() as db: task = None try: result = await db.execute( select(ProjectVectorizationTask).where( ProjectVectorizationTask.task_id == task_id ) ) task = result.scalar_one_or_none() if not task: return project_result = await db.execute( select(Project).where(Project.id == task.project_id) ) project = project_result.scalar_one_or_none() if not project: task.status = "failed" task.error_message = "项目不存在" task.finished_at = func.now() await db.commit() return task.status = "running" task.started_at = func.now() await db.commit() result_data = await zvec_service.vectorize_project( db, task.project_id, project.storage_key, force=task.task_type == "full", ) task.status = "success" if result_data.get("failed", 0) == 0 else "failed" task.total = result_data.get("total", 0) task.processed = result_data.get("processed", 0) task.skipped = result_data.get("skipped", 0) task.failed = result_data.get("failed", 0) task.error_message = ( f"{task.failed} 个文件向量化失败" if task.failed else None ) task.finished_at = func.now() await db.commit() except Exception as exc: logger.exception("Project vectorization task failed: %s", task_id) await db.rollback() if task is not None: task.status = "failed" task.error_message = str(exc)[:2000] task.finished_at = func.now() await db.commit() project_vectorization_task_service = ProjectVectorizationTaskService()