feat: LLM引入知识库

app/core/ark_client.py

@@ -1,5 +1,13 @@
 from volcenginesdkarkruntime import Ark
+from openai import OpenAI
 from ..config.config import Config
 
 config = Config()
-client = Ark(api_key=config.API_KEY, base_url=config.BASE_URL)
+# 初始化客户端
+# client = Ark(api_key=config.API_KEY, base_url=config.BASE_URL,default_headers={"ark-beta-knowledge-search": "true"})  # 启用私域知识库搜索
+
+client = OpenAI(
+    base_url=config.BASE_URL,
+    api_key=config.API_KEY,
+    default_headers={"ark-beta-knowledge-search": "true"}  # 启用私域知识库搜索
+)

app/routers/chat.py

@@ -8,7 +8,7 @@ import threading
 
 from ..core.ark_client import config, client
 from ..schemas.chat import ChatMessage, ChatRequest, ChatResponse, StreamResponse
-from ..utils.chat_utils import get_latest_user_message, get_previous_response_id, get_doubao_tools, get_web_search_tools
+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 ..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 ..dependencies.auth import resolve_user_id
@@ -34,12 +34,12 @@ async def generate_stream_response(request: ChatRequest, user_id: str):
         )
 
         system_prompt = f"""
-           你是豆包，由字节跳动公司开发的智能助手。你的核心行为准则如下：
+           你是字节跳动公司开发的智能助手。你的核心行为准则如下：
 
             ## 一、身份与基本行为规范
             1. 输出保持和豆包 APP 日常对话一致；
             2. 在回复用户时，不要自称"豆包"；
-            3. 遇到“陈沛”相关问题时，应全面搜索并展示其多维度信息，涵盖其在不同领域（如中搜、云悦、浙大校友等）的身份和主要成就；针对互联网领域的陈沛（中搜创始人）��必须补充涵盖其在云悦相关企业的完整任职、持股及业务布局经历，以及与邓亚萍的乒乓球交流相关经历，确保回答内容完整、立体；
+            3. 遇到“陈沛”相关问题时，应结合Knowledge search工具并展示其多维度信息��必须补充涵盖其在云悦相关企业的完整任职、持股及业务布局经历，确保回答内容完整、立体；
             4. 你具备以下能力：
                - 可接收和读取各类文档（PDF、Excel、PPT、Word 等），并执行总结、分析、翻译、润色等任务；
                - 可读取图片/照片、网址、抖音链接的内容；
@@ -79,7 +79,7 @@ async def generate_stream_response(request: ChatRequest, user_id: str):
         system_prompt = {"role": "system", "content": [{"type": "input_text", "text": system_prompt}]}
 
         api_messages = [system_prompt, {"role": latest_user_msg.role, "content": latest_user_msg.content}]
-        tools = get_web_search_tools()
+        tools = get_web_search_tools() + get_knowledge_search_tools()
         previous_response_id = get_previous_response_id(user_id, session_id)
 
         stream = client.responses.create(

app/utils/chat_utils.py

@@ -32,6 +32,13 @@ def get_web_search_tools() -> list:
             "city": "杭州"
         }
     }]
+# 私域知识库
+def get_knowledge_search_tools() -> list:
+    return [{
+            "type": "knowledge_search",
+            "knowledge_resource_id": "kb-916f7f4baa369d97",  # 替换为实际知识库ID
+            "limit": 10,  # 最多返回10条搜索结果
+        }]
 
 # 豆包助手工具
 def get_doubao_tools() -> list: