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 字符串。
     保存知识库文档上传记录,返回插入的 ObjectId 字符串。
     doc_data 应包含:doc_id, resource_id, collection_name, project,
     doc_data 应包含:doc_id, resource_id, collection_name, project,
-                     tos_path, doc_name, doc_type, app_name
+                     doc_name, doc_type, app_name
     """
     """
     try:
     try:
         _ensure_knowledge_docs_index()
         _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(
         docs = list(
             knowledge_docs.find(query, {"_id": 1, "doc_id": 1, "resource_id": 1,
             knowledge_docs.find(query, {"_id": 1, "doc_id": 1, "resource_id": 1,
                                         "collection_name": 1, "project": 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})
                                         "app_name": 1, "created_at": 1})
             .sort("created_at", -1)
             .sort("created_at", -1)
             .skip(skip)
             .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:
     except Exception as e:
         print(f"MongoDB 知识库文档记录查询失败: {e}")
         print(f"MongoDB 知识库文档记录查询失败: {e}")
         return {"total": 0, "page": page, "page_size": page_size, "list": []}
         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 os
 import json
 import json
 import uuid
 import uuid
+import asyncio
 import tempfile
 import tempfile
 from typing import Optional
 from typing import Optional
 
 
 from fastapi import APIRouter, HTTPException, UploadFile, File, Form
 from fastapi import APIRouter, HTTPException, UploadFile, File, Form
 from pydantic import BaseModel
 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.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()
 router = APIRouter()
 
 
@@ -24,6 +25,11 @@ class KnowledgeDocsQuery(BaseModel):
     page: int = 1
     page: int = 1
     page_size: int = 20
     page_size: int = 20
 
 
+
+class KnowledgeDeleteRequest(BaseModel):
+    doc_id: str
+    app_name: str
+
 # 支持的文档类型:后缀 → Viking doc_type 值
 # 支持的文档类型:后缀 → Viking doc_type 值
 ALLOWED_DOC_EXTENSIONS: dict[str, str] = {
 ALLOWED_DOC_EXTENSIONS: dict[str, str] = {
     ".txt":      "txt",
     ".txt":      "txt",
@@ -45,6 +51,14 @@ def _generate_doc_id() -> str:
     return f"doc_{uuid.uuid4().hex}"
     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="上传本地文档到知识库")
 @router.post("/upload", summary="上传本地文档到知识库")
 async def upload_knowledge_doc(
 async def upload_knowledge_doc(
     file: UploadFile = File(..., description="待上传文档,支持 txt/doc/docx/pdf/markdown/pptx/ppt"),
     file: UploadFile = File(..., description="待上传文档,支持 txt/doc/docx/pdf/markdown/pptx/ppt"),
@@ -152,18 +166,20 @@ async def upload_knowledge_doc(
         except Exception as e:
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"导入知识库失败: {str(e)}")
             raise HTTPException(status_code=500, detail=f"导入知识库失败: {str(e)}")
 
 
+        # ── 9. 异步删除 TOS 临时文件(不阻塞响应)─────────────────
+        asyncio.create_task(_cleanup_tos(object_key))
+
     finally:
     finally:
-        # ── 9. 清理临时文件 ──────────────────────────────────────
+        # ── 10. 清理本地临时文件 ─────────────────────────────────
         if tmp_path and os.path.exists(tmp_path):
         if tmp_path and os.path.exists(tmp_path):
             os.unlink(tmp_path)
             os.unlink(tmp_path)
 
 
-    # ── 10. 构造返回结果并存储到 MongoDB ──────────────────────────
+    # ── 11. 构造返回结果并存储到 MongoDB ─────────────────────────
     result = {
     result = {
         "doc_id": kb_data.get("doc_id", final_doc_id),
         "doc_id": kb_data.get("doc_id", final_doc_id),
         "resource_id": kb_data.get("resource_id", resource_id),
         "resource_id": kb_data.get("resource_id", resource_id),
         "collection_name": kb_data.get("collection_name", ""),
         "collection_name": kb_data.get("collection_name", ""),
         "project": kb_data.get("project", "default"),
         "project": kb_data.get("project", "default"),
-        "tos_path": tos_uri,
         "doc_name": final_doc_name,
         "doc_name": final_doc_name,
         "doc_type": doc_type,
         "doc_type": doc_type,
         "app_name": app_name,
         "app_name": app_name,
@@ -182,4 +198,37 @@ async def list_knowledge_docs(request: KnowledgeDocsQuery):
     """
     """
     page = max(1, request.page)
     page = max(1, request.page)
     page_size = request.page_size if 1 <= request.page_size <= 100 else 20
     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": "删除成功"}