""" 知识库工具路由 提供本地文档上传到 Viking 知识库的接口，无需 JWT 鉴权。 接口：POST /main/knowledgeTool/upload """ import os import json import uuid import asyncio import tempfile from typing import Optional from fastapi import APIRouter, HTTPException, UploadFile, File, Form from pydantic import BaseModel from ..core.viking_client import upload_to_tos, add_doc_to_knowledge, delete_from_tos, delete_doc_from_knowledge from ..db.ai_config import get_config_by_app_name from ..db.mongo import save_knowledge_doc, get_knowledge_docs, get_knowledge_doc_by_doc_id, delete_knowledge_doc router = APIRouter() class KnowledgeDocsQuery(BaseModel): app_name: str page: int = 1 page_size: int = 20 class KnowledgeDeleteRequest(BaseModel): doc_id: str app_name: str # 支持的文档类型：后缀 → Viking doc_type 值 ALLOWED_DOC_EXTENSIONS: dict[str, str] = { ".txt": "txt", ".doc": "doc", ".docx": "docx", ".pdf": "pdf", ".md": "markdown", ".markdown": "markdown", ".pptx": "pptx", ".ppt": "ppt", } # 单文件大小上限：50 MB MAX_DOC_SIZE = 50 * 1024 * 1024 def _generate_doc_id() -> str: """生成符合 Viking 命名规范的唯一 doc_id（字母开头，仅含字母/数字/下划线）""" return f"doc_{uuid.uuid4().hex}" async def _cleanup_tos(object_key: str) -> None: """上传成功后异步清理 TOS 临时文件，失败仅打印日志不影响主流程。""" try: await delete_from_tos(object_key) except Exception as e: print(f"TOS 临时文件清理失败（key={object_key}）: {e}") @router.post("/upload", summary="上传本地文档到知识库") async def upload_knowledge_doc( file: UploadFile = File(..., description="待上传文档，支持 txt/doc/docx/pdf/markdown/pptx/ppt"), app_name: str = Form(..., description="应用包名，用于查询对应知识库配置"), doc_id: Optional[str] = Form(None, description="文档唯一标识，不填则自动生成"), doc_name: Optional[str] = Form(None, description="文档显示名称，不填则使用原始文件名"), path_segments: Optional[str] = Form(None, description="目录路径 JSON 数组，如 [\"产品文档\",\"V2.0\"]"), tag_list: Optional[str] = Form(None, description="标签列表 JSON 数组，格式与 Viking API 一致"), ): """ 上传本地文档到火山引擎 Viking 知识库。 流程： 1. 校验文件类型和大小 2. 从 ai_config 查询 knowledgeId（作为 resource_id） 3. 上传文件到 TOS 对象存储 4. 调用 Viking /api/knowledge/doc/v2/add 导入知识库 5. 返回 doc_id 和知识库信息 """ # ── 1. 校验文件类型 ────────────────────────────────────────── original_filename = file.filename or "" _, ext = os.path.splitext(original_filename.lower()) if ext not in ALLOWED_DOC_EXTENSIONS: allowed = "、".join(ALLOWED_DOC_EXTENSIONS.keys()) raise HTTPException( status_code=400, detail=f"不支持的文件类型 '{ext}'，仅支持：{allowed}", ) doc_type = ALLOWED_DOC_EXTENSIONS[ext] # ── 2. 读取并校验文件大小 ──────────────────────────────────── file_bytes = await file.read() if len(file_bytes) > MAX_DOC_SIZE: raise HTTPException( status_code=400, detail=f"文件大小超过限制（最大 50 MB，当前 {len(file_bytes) / 1024 / 1024:.1f} MB）", ) if len(file_bytes) == 0: raise HTTPException(status_code=400, detail="文件内容为空") # ── 3. 查询知识库配置 ──────────────────────────────────────── ai_config = get_config_by_app_name(app_name) if not ai_config: raise HTTPException(status_code=404, detail=f"未找到 appName '{app_name}' 的配置") resource_id = ai_config.get("knowledgeId", "") if not resource_id: raise HTTPException( status_code=400, detail=f"appName '{app_name}' 的配置中未设置 knowledgeId", ) # ── 4. 解析可选 JSON 参数 ──────────────────────────────────── # 强制携带 app_name 标签，前端额外传入的标签追加在后面 parsed_tag_list: list = [ {"field_name": "app_name", "field_type": "string", "field_value": app_name} ] if tag_list: try: extra_tags = json.loads(tag_list) parsed_tag_list.extend(extra_tags) except json.JSONDecodeError: raise HTTPException(status_code=400, detail="tag_list 格式错误，需为合法的 JSON 数组") parsed_path_segments: list | None = None if path_segments: try: parsed_path_segments = json.loads(path_segments) except json.JSONDecodeError: raise HTTPException(status_code=400, detail="path_segments 格式错误，需为合法的 JSON 数组") # ── 5. 确定 doc_id 和 doc_name ────────────────────────────── final_doc_id = doc_id.strip() if doc_id else _generate_doc_id() final_doc_name = doc_name.strip() if doc_name else original_filename # ── 6. 写临时文件 ──────────────────────────────────────────── tmp_path = None try: suffix = ext # 保留原后缀，TOS SDK 可能依赖文件名做 Content-Type 推断 with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: tmp.write(file_bytes) tmp_path = tmp.name # ── 7. 上传到 TOS ──────────────────────────────────────── try: object_key, tos_uri = await upload_to_tos(tmp_path, app_name, original_filename) except RuntimeError as e: raise HTTPException(status_code=500, detail=str(e)) except Exception as e: raise HTTPException(status_code=500, detail=f"TOS 上传失败: {str(e)}") # ── 8. 调用 Viking API 导入知识库 ─────────────────────── try: kb_data = await add_doc_to_knowledge( resource_id=resource_id, doc_id=final_doc_id, tos_uri=tos_uri, doc_name=final_doc_name, doc_type=doc_type, tag_list=parsed_tag_list, path_segments=parsed_path_segments, ) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) except Exception as e: raise HTTPException(status_code=500, detail=f"导入知识库失败: {str(e)}") # ── 9. 异步删除 TOS 临时文件（不阻塞响应）───────────────── asyncio.create_task(_cleanup_tos(object_key)) finally: # ── 10. 清理本地临时文件 ───────────────────────────────── if tmp_path and os.path.exists(tmp_path): os.unlink(tmp_path) # ── 11. 构造返回结果并存储到 MongoDB ───────────────────────── result = { "doc_id": kb_data.get("doc_id", final_doc_id), "resource_id": kb_data.get("resource_id", resource_id), "collection_name": kb_data.get("collection_name", ""), "project": kb_data.get("project", "default"), "doc_name": final_doc_name, "doc_type": doc_type, "app_name": app_name, } inserted_id = save_knowledge_doc(result.copy()) result["id"] = inserted_id or "" return result @router.post("/docs", summary="分页查询知识库文档上传记录") async def list_knowledge_docs(request: KnowledgeDocsQuery): """ 按 app_name 分页查询已上传的知识库文档记录，按上传时间倒序。 请求体使用 JSON，方便后续扩展查询条件。 """ page = max(1, request.page) page_size = request.page_size if 1 <= request.page_size <= 100 else 20 return get_knowledge_docs(app_name=request.app_name, page=page, page_size=page_size) @router.post("/delete", summary="删除知识库文档") async def delete_knowledge_doc_api(request: KnowledgeDeleteRequest): """ 删除指定文档，依次执行： 1. 按 doc_id 查询 MongoDB，获取 resource_id 2. 调用 Viking /api/knowledge/doc/delete 删除知识库文档 3. 删除 MongoDB 中的上传记录 注：TOS 临时文件由桶生命周期规则自动过期清理，此处不处理。 """ # ── 1. 查询 MongoDB 记录 ──────────────────────────────────── doc = get_knowledge_doc_by_doc_id(request.doc_id) if not doc: raise HTTPException(status_code=404, detail=f"未找到 doc_id '{request.doc_id}' 的文档记录") if doc.get("app_name") != request.app_name: raise HTTPException(status_code=403, detail="doc_id 与 app_name 不匹配") resource_id = doc.get("resource_id", "") # ── 2. 删除 Viking 知识库文档 ─────────────────────────────── try: await delete_doc_from_knowledge(resource_id=resource_id, doc_id=request.doc_id) except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) except Exception as e: raise HTTPException(status_code=500, detail=f"删除 Viking 知识库文档失败: {str(e)}") # ── 3. 删除 MongoDB 记录 ──────────────────────────────────── delete_knowledge_doc(request.doc_id) return {"doc_id": request.doc_id, "message": "删除成功"}