Browse Source

feat: - viking知识库-本地上传成功后同时删除tos内的文件
- 删除viking知识库文档

zhangwl 1 week ago
parent
commit
bf580d8fc4
3 changed files with 245 additions and 8 deletions
  1. 160 0
      app/core/viking_client.py
  2. 30 2
      app/db/mongo.py
  3. 55 6
      app/routers/knowledge_tool.py

+ 160 - 0
app/core/viking_client.py

@@ -0,0 +1,160 @@
+"""
+Viking 知识库上传工具模块
+职责:
+  1. 将本地临时文件上传到火山引擎 TOS 对象存储
+  2. 调用 Viking /api/knowledge/doc/v2/add 接口将 TOS 文件导入知识库
+"""
+import os
+import uuid
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+
+import httpx
+
+from ..config.config import Config
+
+cfg = Config()
+
+# 线程池用于执行 TOS 同步 SDK 调用
+_executor = ThreadPoolExecutor(max_workers=4)
+
+
+def _tos_upload_sync(local_path: str, object_key: str) -> str:
+    """
+    同步上传本地文件到 TOS,返回 tos:// URI。
+    在线程池中调用,避免阻塞事件循环。
+    """
+    try:
+        import tos as tos_sdk
+    except ImportError:
+        raise RuntimeError("缺少 tos 依赖,请执行 pip install tos")
+
+    client = tos_sdk.TosClientV2(
+        ak=cfg.VIKING_AK,
+        sk=cfg.VIKING_SK,
+        endpoint=cfg.VIKING_TOS_ENDPOINT,
+        region=cfg.VIKING_TOS_REGION,
+    )
+    client.put_object_from_file(
+        bucket=cfg.VIKING_TOS_BUCKET,
+        key=object_key,
+        file_path=local_path,
+    )
+    return f"tos://{cfg.VIKING_TOS_BUCKET}/{object_key}"
+
+
+async def upload_to_tos(local_path: str, app_name: str, original_filename: str) -> tuple[str, str]:
+    """
+    异步上传本地文件到 TOS。
+    返回 (object_key, tos_uri)。
+    object_key 格式:{app_name}/{uuid32}_{原始文件名}
+    """
+    # 清理文件名中的特殊字符,只保留字母数字下划线中文和常见标点
+    safe_name = os.path.basename(original_filename)
+    uid = uuid.uuid4().hex
+    object_key = f"{app_name}/{uid}_{safe_name}"
+
+    loop = asyncio.get_event_loop()
+    tos_uri = await loop.run_in_executor(_executor, _tos_upload_sync, local_path, object_key)
+    return object_key, tos_uri
+
+
+async def add_doc_to_knowledge(
+    resource_id: str,
+    doc_id: str,
+    tos_uri: str,
+    doc_name: str,
+    doc_type: str,
+    tag_list: list | None = None,
+    path_segments: list | None = None,
+) -> dict:
+    """
+    调用 Viking /api/knowledge/doc/v2/add 接口,将 TOS 文件导入知识库。
+    使用 resource_id 指定知识库(对应 ai_config.knowledgeId)。
+    返回响应体 data 字段。
+    """
+    url = f"https://{cfg.VIKING_HOST}/api/knowledge/doc/v2/add"
+    payload: dict = {
+        "resource_id": resource_id,
+        "doc_id": doc_id,
+        "doc_name": doc_name,
+        "doc_type": doc_type,
+        "uri": tos_uri,
+    }
+    if tag_list:
+        payload["tag_list"] = tag_list
+    if path_segments:
+        payload["path_segments"] = path_segments
+
+    async with httpx.AsyncClient(timeout=30) as client:
+        resp = await client.post(
+            url,
+            json=payload,
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {cfg.VIKING_API_KEY}",
+            },
+        )
+        body = resp.json()
+        code = body.get("code", -1)
+        if code != 0:
+            raise ValueError(
+                f"Viking API 返回错误: {body.get('message')} "
+                f"(code={code}, request_id={body.get('request_id')})"
+            )
+        return body["data"]
+
+
+def _tos_delete_sync(object_key: str) -> None:
+    """
+    同步删除 TOS 对象。
+    在线程池中调用,避免阻塞事件循环。
+    """
+    try:
+        import tos as tos_sdk
+    except ImportError:
+        raise RuntimeError("缺少 tos 依赖,请执行 pip install tos")
+
+    client = tos_sdk.TosClientV2(
+        ak=cfg.VIKING_AK,
+        sk=cfg.VIKING_SK,
+        endpoint=cfg.VIKING_TOS_ENDPOINT,
+        region=cfg.VIKING_TOS_REGION,
+    )
+    client.delete_object(bucket=cfg.VIKING_TOS_BUCKET, key=object_key)
+
+
+async def delete_from_tos(object_key: str) -> None:
+    """异步删除 TOS 对象。"""
+    loop = asyncio.get_event_loop()
+    await loop.run_in_executor(_executor, _tos_delete_sync, object_key)
+
+
+async def delete_doc_from_knowledge(resource_id: str, doc_id: str) -> dict:
+    """
+    调用 Viking /api/knowledge/doc/delete 接口删除知识库文档。
+    返回完整响应体。
+    """
+    url = f"https://{cfg.VIKING_HOST}/api/knowledge/doc/delete"
+    payload = {
+        "resource_id": resource_id,
+        "doc_id": doc_id,
+    }
+
+    async with httpx.AsyncClient(timeout=30) as client:
+        resp = await client.post(
+            url,
+            json=payload,
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {cfg.VIKING_API_KEY}",
+            },
+        )
+        body = resp.json()
+        code = body.get("code", -1)
+        if code != 0:
+            raise ValueError(
+                f"Viking API 返回错误: {body.get('message')} "
+                f"(code={code}, request_id={body.get('request_id')})"
+            )
+        return body

+ 30 - 2
app/db/mongo.py

@@ -371,7 +371,7 @@ def save_knowledge_doc(doc_data: dict) -> str | None:
     """
     保存知识库文档上传记录,返回插入的 ObjectId 字符串。
     doc_data 应包含:doc_id, resource_id, collection_name, project,
-                     tos_path, doc_name, doc_type, app_name
+                     doc_name, doc_type, app_name
     """
     try:
         _ensure_knowledge_docs_index()
@@ -396,7 +396,7 @@ def get_knowledge_docs(app_name: str, page: int = 1, page_size: int = 20) -> dic
         docs = list(
             knowledge_docs.find(query, {"_id": 1, "doc_id": 1, "resource_id": 1,
                                         "collection_name": 1, "project": 1,
-                                        "tos_path": 1, "doc_name": 1, "doc_type": 1,
+                                        "doc_name": 1, "doc_type": 1,
                                         "app_name": 1, "created_at": 1})
             .sort("created_at", -1)
             .skip(skip)
@@ -410,3 +410,31 @@ def get_knowledge_docs(app_name: str, page: int = 1, page_size: int = 20) -> dic
     except Exception as e:
         print(f"MongoDB 知识库文档记录查询失败: {e}")
         return {"total": 0, "page": page, "page_size": page_size, "list": []}
+
+
+def get_knowledge_doc_by_doc_id(doc_id: str) -> dict | None:
+    """
+    按 doc_id 查询单条知识库文档记录。
+    返回文档 dict(含 id 字符串字段),未找到返回 None。
+    """
+    try:
+        doc = knowledge_docs.find_one({"doc_id": doc_id})
+        if doc:
+            doc["id"] = str(doc.pop("_id"))
+        return doc
+    except Exception as e:
+        print(f"MongoDB 知识库文档记录查询失败: {e}")
+        return None
+
+
+def delete_knowledge_doc(doc_id: str) -> bool:
+    """
+    按 doc_id 删除知识库文档记录。
+    返回 True 表示成功删除,False 表示记录不存在或删除异常。
+    """
+    try:
+        result = knowledge_docs.delete_one({"doc_id": doc_id})
+        return result.deleted_count > 0
+    except Exception as e:
+        print(f"MongoDB 知识库文档记录删除失败: {e}")
+        return False

+ 55 - 6
app/routers/knowledge_tool.py

@@ -6,15 +6,16 @@
 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
+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
+from ..db.mongo import save_knowledge_doc, get_knowledge_docs, get_knowledge_doc_by_doc_id, delete_knowledge_doc
 
 router = APIRouter()
 
@@ -24,6 +25,11 @@ class KnowledgeDocsQuery(BaseModel):
     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",
@@ -45,6 +51,14 @@ def _generate_doc_id() -> str:
     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"),
@@ -152,18 +166,20 @@ async def upload_knowledge_doc(
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"导入知识库失败: {str(e)}")
 
+        # ── 9. 异步删除 TOS 临时文件(不阻塞响应)─────────────────
+        asyncio.create_task(_cleanup_tos(object_key))
+
     finally:
-        # ── 9. 清理临时文件 ──────────────────────────────────────
+        # ── 10. 清理本地临时文件 ─────────────────────────────────
         if tmp_path and os.path.exists(tmp_path):
             os.unlink(tmp_path)
 
-    # ── 10. 构造返回结果并存储到 MongoDB ──────────────────────────
+    # ── 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"),
-        "tos_path": tos_uri,
         "doc_name": final_doc_name,
         "doc_type": doc_type,
         "app_name": app_name,
@@ -182,4 +198,37 @@ async def list_knowledge_docs(request: KnowledgeDocsQuery):
     """
     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)
+    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": "删除成功"}