nex_docus/backend/app/services/project_vectorization_task_service.py

"""
项目向量化后台任务服务
"""
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()