Browse Source

feat: - 自建知识库

zhangwl 4 weeks ago
parent
commit
91a9a8353b
4 changed files with 65 additions and 5 deletions
  1. 26 2
      app/routers/chat.py
  2. 32 2
      app/routers/moderation.py
  3. 2 1
      main.py
  4. 5 0
      requirements.txt

+ 26 - 2
app/routers/chat.py

@@ -13,6 +13,7 @@ from ..core.asr_client import transcribe_audio
 from ..config.config import Config
 from ..schemas.chat import ChatMessage, ChatRequest, ChatResponse, StreamResponse, FileAttachment
 from ..utils.chat_utils import get_latest_user_message, get_previous_response_id, get_doubao_tools, get_web_search_tools, get_knowledge_search_tools
+from ..utils.kb_search import search_knowledge_base, format_knowledge_context
 from ..routers.users import get_current_active_user, User
 from ..db.mongo import save_chat_log, save_chat_history, get_chat_history, delete_chat_history, get_sessions
 from ..db.ai_config import get_config_by_app_name
@@ -155,7 +156,8 @@ async def generate_stream_response(request: ChatRequest, user_context: UserConte
             raise ValueError(f"未找到appName '{user_context.app_name}' 的配置")
 
         client = get_client(user_context.app_name)
-        knowledge_id = ai_config["knowledgeId"]
+        knowledge_id = ai_config.get("knowledgeId", "")  # 豆包原生知识库ID(可选)
+        use_custom_kb = ai_config.get("use_custom_kb", True)  # 是否使用自建知识库,默认启用
 
         latest_user_msg = get_latest_user_message(request.messages)
         if not latest_user_msg:
@@ -172,6 +174,21 @@ async def generate_stream_response(request: ChatRequest, user_context: UserConte
             attachments=user_attachments,
         )
 
+        # 自建知识库检索(前置到prompt)
+        kb_context = ""
+        if use_custom_kb:
+            try:
+                kb_results = search_knowledge_base(
+                    query=latest_user_msg.content,
+                    tenant_id=user_context.app_name,
+                    top_k=5
+                )
+                if kb_results:
+                    kb_context = format_knowledge_context(kb_results)
+                    print(f"知识库检索成功,命中 {len(kb_results)} 个结果")
+            except Exception as e:
+                print(f"知识库检索异常: {str(e)}")
+
         system_prompt = f"""
            你是云悦公司开发的智能助手。你的核心行为准则如下:
 
@@ -300,11 +317,18 @@ async def generate_stream_response(request: ChatRequest, user_context: UserConte
             - 用结构化语言给出准确、稳妥、不过度暴露内部过程的回答。
 
            """
+
+        # 如果检索到知识库内容,追加到 system prompt
+        if kb_context:
+            system_prompt += f"\n\n{kb_context}\n"
+
         system_prompt = {"role": "system", "content": [{"type": "input_text", "text": system_prompt}]}
 
         api_messages = [system_prompt, build_multimodal_input(latest_user_msg)]
         tools = get_web_search_tools()
-        if knowledge_id:
+
+        # 如果使用豆包原生知识库(兼容旧配置)
+        if knowledge_id and not use_custom_kb:
             tools += get_knowledge_search_tools(knowledge_id)
         previous_response_id = get_previous_response_id(user_context.user_id, session_id)
 

File diff suppressed because it is too large
+ 32 - 2
app/routers/moderation.py


+ 2 - 1
main.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI
 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_base
 
 # 创建FastAPI应用实例
 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(ai_config.router, prefix="/main/aiConfig", tags=["AI配置管理"])
 app.include_router(moderation.router, prefix="/main/moderation", tags=["内容审核"])
+app.include_router(knowledge_base.router, prefix="/main/kb", tags=["知识库管理"])
 
 
 

+ 5 - 0
requirements.txt

@@ -13,3 +13,8 @@ httpx==0.28.1
 openai==2.30.0
 python-multipart==0.0.32
 websockets>=12.0
+qdrant-client==1.7.0
+oss2==2.18.0
+python-docx==1.1.0
+openpyxl==3.1.2
+PyPDF2==3.0.1