Browse Source

feat: - viking知识库-支持本地上传文本到知识库和分页查询

zhangwl 1 week ago
parent
commit
1d688cd467
6 changed files with 267 additions and 3 deletions
  1. 13 2
      .env
  2. 9 0
      app/config/config.py
  3. 57 0
      app/db/mongo.py
  4. 185 0
      app/routers/knowledge_tool.py
  5. 2 1
      main.py
  6. 1 0
      requirements.txt

+ 13 - 2
.env

@@ -1,8 +1,7 @@
 # ark的key , 也就是火山引擎的key - 云享生态,后续需要改成根据包名进行映射
 # ark的key , 也就是火山引擎的key - 云享生态,后续需要改成根据包名进行映射
 ARK_API_KEY=ark-95d3c2b6-b694-46b8-a3b2-5b573e1596ca-cf7e4
 ARK_API_KEY=ark-95d3c2b6-b694-46b8-a3b2-5b573e1596ca-cf7e4
 ARK_BASE_URL=https://ark.cn-beijing.volces.com/api/v3
 ARK_BASE_URL=https://ark.cn-beijing.volces.com/api/v3
-# 知识库的api key
-ARK_KNOW_API_KEY=W1D0N6X3GAQTFJXWAYAV09CPPMRZ8VP01YRXE028HTB2ZZPDBC4G60R30C9P6CVK8
+
 
 
 # 豆包的2.0模型 doubao-seed-2-0-lite-260215
 # 豆包的2.0模型 doubao-seed-2-0-lite-260215
 #MODEL_NAME=ep-20260430173618-626ng
 #MODEL_NAME=ep-20260430173618-626ng
@@ -36,3 +35,15 @@ ENVIRONMENT=production
 TENCENT_IM_APP_ID=1400017757
 TENCENT_IM_APP_ID=1400017757
 TENCENT_IM_ADMIN_ID=adminjava
 TENCENT_IM_ADMIN_ID=adminjava
 TENCENT_IM_SECRET_KEY=f8a4c34728e74d235da1ea2b1775a39c27787a4954e379de5ac424806e8b67e3
 TENCENT_IM_SECRET_KEY=f8a4c34728e74d235da1ea2b1775a39c27787a4954e379de5ac424806e8b67e3
+
+
+# Viking 知识库上传配置
+VIKING_API_KEY=W1D0N6X3GAQTFJXWAYAV09CPPMRZ8VP01YRXE028HTB2ZZPDBC4G60R30C9P6CVK8
+# 火山引擎-子账号-ak & sk
+VIKING_AK=AKLTNDE4YmE3MjU0MjE3NDIwYzgwZmNlZjA1MDRlNDNiNzM
+VIKING_SK=TlRaaU5HVXdNbVJoT0dNeU5EZ3pOVGs1Tm1KaU9UazVOekppTlRSbE0yVQ==
+
+VIKING_TOS_BUCKET=auto-knowledge-doc
+VIKING_TOS_ENDPOINT=tos-cn-beijing.volces.com
+VIKING_TOS_REGION=cn-beijing
+VIKING_HOST=api-knowledgebase.mlp.cn-beijing.volces.com

+ 9 - 0
app/config/config.py

@@ -18,3 +18,12 @@ class Config:
     # 火山引擎 ASR 语音识别配置
     # 火山引擎 ASR 语音识别配置
     VOLC_ASR_API_KEY: str = os.getenv("VOLC_ASR_API_KEY", "")
     VOLC_ASR_API_KEY: str = os.getenv("VOLC_ASR_API_KEY", "")
     VOLC_ASR_RESOURCE_ID: str = os.getenv("VOLC_ASR_RESOURCE_ID", "volc.seedasr.sauc.duration")
     VOLC_ASR_RESOURCE_ID: str = os.getenv("VOLC_ASR_RESOURCE_ID", "volc.seedasr.sauc.duration")
+
+    # Viking 知识库上传配置
+    VIKING_API_KEY: str = os.getenv("VIKING_API_KEY", "")
+    VIKING_AK: str = os.getenv("VIKING_AK", "")
+    VIKING_SK: str = os.getenv("VIKING_SK", "")
+    VIKING_TOS_BUCKET: str = os.getenv("VIKING_TOS_BUCKET", "auto-knowledge-doc")
+    VIKING_TOS_ENDPOINT: str = os.getenv("VIKING_TOS_ENDPOINT", "tos-cn-beijing.volces.com")
+    VIKING_TOS_REGION: str = os.getenv("VIKING_TOS_REGION", "cn-beijing")
+    VIKING_HOST: str = os.getenv("VIKING_HOST", "api-knowledgebase.mlp.cn-beijing.volces.com")

+ 57 - 0
app/db/mongo.py

@@ -20,6 +20,8 @@ chat_history_col = db["chat_history"]
 circle_prompts = db["circle_prompt"]
 circle_prompts = db["circle_prompt"]
 # 历史人物集合
 # 历史人物集合
 historical_figures = db["historical_figures"]
 historical_figures = db["historical_figures"]
+# 知识库文档上传记录
+knowledge_docs = db["knowledge_docs"]
 
 
 def _ensure_index():
 def _ensure_index():
     try:
     try:
@@ -353,3 +355,58 @@ def get_group_ai_reply_count(session_id: str, ai_role: str = None) -> int:
         print(f"MongoDB 查询AI回复次数失败: {e}")
         print(f"MongoDB 查询AI回复次数失败: {e}")
         return 0
         return 0
 
 
+
+
+# ==================== 知识库文档上传记录 ====================
+
+def _ensure_knowledge_docs_index():
+    try:
+        knowledge_docs.create_index([("app_name", 1), ("created_at", -1)])
+        knowledge_docs.create_index([("doc_id", 1)], unique=True)
+    except Exception:
+        pass
+
+
+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
+    """
+    try:
+        _ensure_knowledge_docs_index()
+        doc_data["created_at"] = datetime.now()
+        result = knowledge_docs.insert_one(doc_data)
+        return str(result.inserted_id)
+    except Exception as e:
+        print(f"MongoDB 知识库文档记录写入失败: {e}")
+        return None
+
+
+def get_knowledge_docs(app_name: str, page: int = 1, page_size: int = 20) -> dict:
+    """
+    分页查询知识库文档上传记录。
+    返回 {"total": int, "page": int, "page_size": int, "list": [...]}
+    """
+    try:
+        _ensure_knowledge_docs_index()
+        query = {"app_name": app_name}
+        total = knowledge_docs.count_documents(query)
+        skip = (page - 1) * page_size
+        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,
+                                        "app_name": 1, "created_at": 1})
+            .sort("created_at", -1)
+            .skip(skip)
+            .limit(page_size)
+        )
+        for doc in docs:
+            doc["id"] = str(doc.pop("_id"))
+            if "created_at" in doc:
+                doc["created_at"] = doc["created_at"].strftime("%Y-%m-%d %H:%M:%S")
+        return {"total": total, "page": page, "page_size": page_size, "list": docs}
+    except Exception as e:
+        print(f"MongoDB 知识库文档记录查询失败: {e}")
+        return {"total": 0, "page": page, "page_size": page_size, "list": []}

+ 185 - 0
app/routers/knowledge_tool.py

@@ -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)

+ 2 - 1
main.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.cors import CORSMiddleware
 
 
-from app.routers import users, chat, chat_tools, ai_config, moderation
+from app.routers import users, chat, chat_tools, ai_config, moderation, knowledge_tool
 
 
 # 创建FastAPI应用实例
 # 创建FastAPI应用实例
 app = FastAPI(title="聊天机器人", version="1.0.0", description="基于fastapi+VUE的聊天机器人")
 app = FastAPI(title="聊天机器人", version="1.0.0", description="基于fastapi+VUE的聊天机器人")
@@ -23,6 +23,7 @@ app.include_router(chat.router, prefix="/main/chat", tags=["聊天管理"])
 app.include_router(chat_tools.router, prefix="/main/chatTools", tags=["AI工具管理"])
 app.include_router(chat_tools.router, prefix="/main/chatTools", tags=["AI工具管理"])
 app.include_router(ai_config.router, prefix="/main/aiConfig", tags=["AI配置管理"])
 app.include_router(ai_config.router, prefix="/main/aiConfig", tags=["AI配置管理"])
 app.include_router(moderation.router, prefix="/main/moderation", tags=["内容审核"])
 app.include_router(moderation.router, prefix="/main/moderation", tags=["内容审核"])
+app.include_router(knowledge_tool.router, prefix="/main/knowledgeTool", tags=["知识库工具"])
 
 
 
 
 
 

+ 1 - 0
requirements.txt

@@ -14,5 +14,6 @@ openai==2.30.0
 tls-sig-api-v2==1.1
 tls-sig-api-v2==1.1
 python-multipart==0.0.32
 python-multipart==0.0.32
 websockets>=12.0
 websockets>=12.0
+tos>=2.7.0
 #llama-index-core==0.14.25
 #llama-index-core==0.14.25
 #Jinja2==3.1.6
 #Jinja2==3.1.6