""" 知识加工 API 路由模块 处理知识库文件的加工任务，包括合并、对比、总结等功能。 """ from typing import Optional import asyncpg from fastapi import APIRouter, Depends, BackgroundTasks, Query from core.dependencies import get_db, get_current_user from core.database import get_db_pool from core.exceptions import NotFoundError, BadRequestError from models.user import User from models.knowledge_processing import ( TaskCreateRequest, TaskResponse, TaskListResponse, TaskStatusResponse, TaskStatus ) from services.knowledge_base_service import KnowledgeBaseService from services.knowledge_processing_service import ( KnowledgeProcessingService, KnowledgeProcessingExecutor ) from utils.helpers import BaseResponse from logger.logging import get_logger logger = get_logger(__name__) # 创建知识加工路由 kb_processing_router = APIRouter(prefix="/api/knowledge-base", tags=["知识加工"]) async def process_task_background(task_id: int): """ 后台任务：执行知识加工 Args: task_id: 任务 ID """ pool = await get_db_pool() async with pool.acquire() as conn: try: logger.info(f"开始后台处理知识加工任务 ID: {task_id}") # 获取任务信息 task = await conn.fetchrow( """ SELECT id, user_id, knowledge_base_id, task_name, instruction, file_ids, task_type, status, result, result_file_url, error_message, created_at, updated_at, started_at, completed_at FROM knowledge_processing_task WHERE id = $1 """, task_id ) if not task: logger.error(f"任务 {task_id} 不存在") return from models.knowledge_processing import KnowledgeProcessingTask task_obj = KnowledgeProcessingTask(**dict(task)) # 更新状态为处理中 await KnowledgeProcessingService.update_task_status( conn, task_id, TaskStatus.PROCESSING ) # 执行任务 success, result, error_message, result_file_url = await KnowledgeProcessingExecutor.process_task( conn, task_obj ) # 更新任务状态 if success: await KnowledgeProcessingService.update_task_status( conn, task_id, TaskStatus.COMPLETED, result=result, result_file_url=result_file_url ) logger.info(f"任务 {task_id} 处理成功，文件链接: {result_file_url}") else: await KnowledgeProcessingService.update_task_status( conn, task_id, TaskStatus.FAILED, error_message=error_message ) logger.error(f"任务 {task_id} 处理失败: {error_message}") except Exception as e: logger.error(f"后台处理任务异常 ID: {task_id}, 错误: {e}") import traceback logger.error(f"错误堆栈: {traceback.format_exc()}") # 更新任务状态为失败 try: await KnowledgeProcessingService.update_task_status( conn, task_id, TaskStatus.FAILED, error_message=str(e) ) except Exception as update_error: logger.error(f"更新任务状态失败: {update_error}") @kb_processing_router.post("/{kb_id}/processing/tasks", response_model=BaseResponse, summary="创建知识加工任务") async def create_processing_task( kb_id: int, task_data: TaskCreateRequest, background_tasks: BackgroundTasks, current_user: User = Depends(get_current_user), conn: asyncpg.Connection = Depends(get_db) ): """ 创建知识加工任务 用户可以选择知识库中的一个或多个文件，输入加工指令，系统将异步处理任务。 支持的任务类型： - merge: 合并文件 - compare: 对比文件 - summary: 总结文件 - custom: 自定义指令 """ try: # 检查知识库是否存在 kb = await KnowledgeBaseService.get_knowledge_base_by_id(conn, kb_id, current_user) if not kb: raise NotFoundError("知识库") # 创建任务 task = await KnowledgeProcessingService.create_task( conn, current_user.id, kb_id, task_data ) # 添加后台处理任务 logger.info(f"添加后台加工任务: task_id={task.id}, type={task.task_type}") background_tasks.add_task(process_task_background, task.id) return BaseResponse( code=200, msg="任务创建成功，正在处理中", data=TaskResponse( id=task.id, task_name=task.task_name, instruction=task.instruction, file_ids=task.file_ids, task_type=task.task_type.value, status=task.status.value, result=task.result, result_file_url=task.result_file_url, error_message=task.error_message, created_at=task.created_at, updated_at=task.updated_at, started_at=task.started_at, completed_at=task.completed_at ).dict() ) except ValueError as e: raise BadRequestError(str(e)) except Exception as e: logger.error(f"创建知识加工任务失败: {e}") raise BadRequestError(f"创建任务失败: {str(e)}") @kb_processing_router.get("/{kb_id}/processing/tasks", response_model=BaseResponse, summary="获取知识加工任务列表") async def get_processing_tasks( kb_id: int, status: Optional[str] = Query(None, description="任务状态筛选"), page: int = Query(1, ge=1, description="页码"), page_size: int = Query(20, ge=1, le=100, description="每页数量"), current_user: User = Depends(get_current_user), conn: asyncpg.Connection = Depends(get_db) ): """获取知识库的加工任务列表""" # 检查知识库是否存在 kb = await KnowledgeBaseService.get_knowledge_base_by_id(conn, kb_id, current_user) if not kb: raise NotFoundError("知识库") # 获取任务列表 tasks, total = await KnowledgeProcessingService.get_user_tasks( conn, current_user.id, kb_id, status, page, page_size ) items = [ TaskResponse( id=task.id, task_name=task.task_name, instruction=task.instruction, file_ids=task.file_ids, task_type=task.task_type.value, status=task.status.value, result=task.result, result_file_url=task.result_file_url, error_message=task.error_message, created_at=task.created_at, updated_at=task.updated_at, started_at=task.started_at, completed_at=task.completed_at ).dict() for task in tasks ] return BaseResponse( code=200, msg="获取任务列表成功", data=TaskListResponse(total=total, items=items).dict() ) @kb_processing_router.get("/{kb_id}/processing/tasks/{task_id}", response_model=BaseResponse, summary="获取任务详情") async def get_task_detail( kb_id: int, task_id: int, current_user: User = Depends(get_current_user), conn: asyncpg.Connection = Depends(get_db) ): """获取知识加工任务详情""" # 检查知识库是否存在 kb = await KnowledgeBaseService.get_knowledge_base_by_id(conn, kb_id, current_user) if not kb: raise NotFoundError("知识库") # 获取任务 task = await KnowledgeProcessingService.get_task_by_id(conn, task_id, current_user.id) if not task or task.knowledge_base_id != kb_id: raise NotFoundError("任务") return BaseResponse( code=200, msg="获取任务详情成功", data=TaskResponse( id=task.id, task_name=task.task_name, instruction=task.instruction, file_ids=task.file_ids, task_type=task.task_type.value, status=task.status.value, result=task.result, result_file_url=task.result_file_url, error_message=task.error_message, created_at=task.created_at, updated_at=task.updated_at, started_at=task.started_at, completed_at=task.completed_at ).dict() ) @kb_processing_router.get("/{kb_id}/processing/tasks/{task_id}/status", response_model=BaseResponse, summary="查询任务处理状态") async def get_task_status( kb_id: int, task_id: int, current_user: User = Depends(get_current_user), conn: asyncpg.Connection = Depends(get_db) ): """ 查询知识加工任务的处理状态（用于前端轮询） Returns: - id: 任务ID - status: pending（待处理）/ processing（处理中）/ completed（已完成）/ failed（失败） - result: 处理结果（仅在completed时返回） - error_message: 错误信息（仅在failed时返回） - updated_at: 更新时间 - started_at: 开始时间 - completed_at: 完成时间 """ # 检查知识库是否存在 kb = await KnowledgeBaseService.get_knowledge_base_by_id(conn, kb_id, current_user) if not kb: raise NotFoundError("知识库") # 获取任务 task = await KnowledgeProcessingService.get_task_by_id(conn, task_id, current_user.id) if not task or task.knowledge_base_id != kb_id: raise NotFoundError("任务") return BaseResponse( code=200, msg="获取任务状态成功", data=TaskStatusResponse( id=task.id, status=task.status.value, result=task.result, result_file_url=task.result_file_url, error_message=task.error_message, updated_at=task.updated_at, started_at=task.started_at, completed_at=task.completed_at ).dict() ) @kb_processing_router.delete("/{kb_id}/processing/tasks/{task_id}", response_model=BaseResponse, summary="删除任务") async def delete_task( kb_id: int, task_id: int, current_user: User = Depends(get_current_user), conn: asyncpg.Connection = Depends(get_db) ): """删除知识加工任务""" # 检查知识库是否存在 kb = await KnowledgeBaseService.get_knowledge_base_by_id(conn, kb_id, current_user) if not kb: raise NotFoundError("知识库") # 获取任务 task = await KnowledgeProcessingService.get_task_by_id(conn, task_id, current_user.id) if not task or task.knowledge_base_id != kb_id: raise NotFoundError("任务") # 删除任务 success = await KnowledgeProcessingService.delete_task(conn, task_id, current_user.id) if not success: raise NotFoundError("任务") return BaseResponse(code=200, msg="删除任务成功", data={"id": task_id})