|
|
@@ -0,0 +1,185 @@
|
|
|
+"""
|
|
|
+知识库工具路由
|
|
|
+提供本地文档上传到 Viking 知识库的接口,无需 JWT 鉴权。
|
|
|
+接口:POST /main/knowledgeTool/upload
|
|
|
+"""
|
|
|
+import os
|
|
|
+import json
|
|
|
+import uuid
|
|
|
+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
|
|
|
+from ..db.ai_config import get_config_by_app_name
|
|
|
+from ..db.mongo import save_knowledge_doc, get_knowledge_docs
|
|
|
+
|
|
|
+router = APIRouter()
|
|
|
+
|
|
|
+
|
|
|
+class KnowledgeDocsQuery(BaseModel):
|
|
|
+ app_name: str
|
|
|
+ page: int = 1
|
|
|
+ page_size: int = 20
|
|
|
+
|
|
|
+# 支持的文档类型:后缀 → 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}"
|
|
|
+
|
|
|
+
|
|
|
+@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)}")
|
|
|
+
|
|
|
+ finally:
|
|
|
+ # ── 9. 清理临时文件 ──────────────────────────────────────
|
|
|
+ if tmp_path and os.path.exists(tmp_path):
|
|
|
+ os.unlink(tmp_path)
|
|
|
+
|
|
|
+ # ── 10. 构造返回结果并存储到 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"),
|
|
|
+ "tos_path": tos_uri,
|
|
|
+ "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)
|