|
@@ -1,234 +0,0 @@
|
|
|
-"""
|
|
|
|
|
-知识库工具路由
|
|
|
|
|
-提供本地文档上传到 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": "删除成功"}
|
|
|