feat: 添加了基于 tavil的web 搜索

.gitignore

@@ -36,3 +36,6 @@ saves
 notebooks
 graphrag
 docker/volumes
+
+
+.cursorrules

requirements.txt

@@ -25,3 +25,4 @@ docx2txt
 uvicorn[standard]
 fastapi
 python-multipart
+tavily-python

src/config/__init__.py

@@ -53,7 +53,7 @@ class Config(SimpleConfig):
         self.add_item("enable_knowledge_base", default=False, des="是否开启知识库")
         self.add_item("enable_knowledge_graph", default=False, des="是否开启知识图谱")
         self.add_item("enable_search_engine", default=False, des="是否开启搜索引擎")
-
+        self.add_item("enable_web_search", default=False, des="是否开启网页搜索")
         # 模型配置
         ## 注意这里是模型名，而不是具体的模型路径，默认使用 HuggingFace 的路径
         ## 如果需要自定义路径，则在 config/base.yaml 中配置 model_local_paths

src/config/models.yaml

@@ -18,6 +18,7 @@ MODEL_NAMES:
       - DEEPSEEK_API_KEY
     models:
       - deepseek-chat
+      - deepseek-reasoner
   zhipu:
     name: 智谱AI (Zhipu)
     url: https://open.bigmodel.cn/dev/api

src/core/startup.py

@@ -1,3 +1,4 @@
+import os
 from src.core import DataBaseManager
 from src.core.retriever import Retriever
 from src.models import select_model
@@ -9,10 +10,29 @@ logger = setup_logger("Startup")
 
 class Startup:
     def __init__(self):
+        self.config = Config("config/base.yaml")
+        self._check_environment()
         self.start()
 
+    def _check_environment(self):
+        """检查必要的环境变量"""
+        required_vars = {
+            "zhipu": ["ZHIPUAI_API_KEY"],
+            "openai": ["OPENAI_API_KEY"],
+            "deepseek": ["DEEPSEEK_API_KEY"],
+        }
+        
+        provider = self.config.model_provider
+        if provider in required_vars:
+            missing = [var for var in required_vars[provider] if not os.getenv(var)]
+            if missing:
+                logger.error(f"Missing required environment variables for {provider}: {missing}")
+                raise ValueError(f"Missing required environment variables: {missing}")
+                
+        if self.config.enable_web_search and not os.getenv("TAVILY_API_KEY"):
+            logger.warning("TAVILY_API_KEY not set, web search will be disabled")
+
     def start(self):
-        self.config = Config()
         self.model = select_model(self.config)
         self.dbm = DataBaseManager(self.config)
         self.retriever = Retriever(self.config, self.dbm, self.model)

src/models/__init__.py

@@ -1,46 +1,23 @@
 from src.utils.logging_config import logger
+from src.models.chat_model import OpenModel, DeepSeek, Zhipu, Qianfan, DashScope, SiliconFlow
+from src.models.embedding import get_embedding_model
 
 
 def select_model(config):
-
+    """
-    model_provider = config.model_provider
+    根据配置选择模型
-    model_name = config.model_name
+    """
-
+    if config.model_provider == "deepseek":
-    logger.info(f"Selecting model from {model_provider} with {model_name}")
+        return DeepSeek(config.model_name)
-
+    elif config.model_provider == "zhipu":
-    if model_provider == "deepseek":
+        return Zhipu(config.model_name)
-        from src.models.chat_model import DeepSeek
+    elif config.model_provider == "openai":
-        return DeepSeek(model_name)
+        return OpenModel(config.model_name)
-
+    elif config.model_provider == "qianfan":
-    elif model_provider == "zhipu":
+        return Qianfan(config.model_name)
-        from src.models.chat_model import Zhipu
+    elif config.model_provider == "dashscope":
-        return Zhipu(model_name)
+        return DashScope(config.model_name)
-
+    elif config.model_provider == "siliconflow":
-    elif model_provider == "qianfan":
+        return SiliconFlow(config.model_name)
-        from src.models.chat_model import Qianfan
-        return Qianfan(model_name)
-
-    elif model_provider == "dashscope":
-        from src.models.chat_model import DashScope
-        return DashScope(model_name)
-
-    elif model_provider == "openai":
-        from src.models.chat_model import OpenModel
-        return OpenModel(model_name)
-
-    elif model_provider == "siliconflow":
-        from src.models.chat_model import SiliconFlow
-        return SiliconFlow(model_name)
-
-    elif model_provider == "custom":
-        model_info = next((x for x in config.custom_models if x["custom_id"] == model_name), None)
-        if model_info is None:
-            raise ValueError(f"Model {model_name} not found in custom models")
-
-        from src.models.chat_model import CustomModel
-        return CustomModel(model_info)
-
-    elif model_provider is None:
-        raise ValueError("Model provider not specified, please modify `model_provider` in `src/config/base.yaml`")
     else:
-        raise ValueError(f"Model provider {model_provider} not supported")
+        raise ValueError(f"Unsupported model provider: {config.model_provider}")

src/models/chat_model.py

@@ -1,6 +1,7 @@
 import os
 from openai import OpenAI
 from src.utils.logging_config import setup_logger
+from zhipuai import ZhipuAI
 
 
 logger = setup_logger(__name__)
@@ -50,8 +51,8 @@ class OpenModel(OpenAIBase):
 class DeepSeek(OpenAIBase):
     def __init__(self, model_name=None):
         model_name = model_name or "deepseek-chat"
-        api_key = os.getenv("DEEPSEEK_API_KEY")
+        api_key = os.getenv("DEEPSEEK_API_KEY", "your-default-api-key")
-        base_url = "https://api.deepseek.com"
+        base_url = os.getenv("DEEPSEEK_API_BASE", "https://api.deepseek.com/v1")
         super().__init__(api_key=api_key, base_url=base_url, model_name=model_name)
 
 
@@ -165,6 +166,14 @@ class DashScope:
         return response.output.choices[0].message
 
 
+class ChatModel:
+    def __init__(self, config):
+        if config.model_provider == "zhipu":
+            self.client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))
+        elif config.model_provider == "openai":
+            self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+
 if __name__ == "__main__":
     model = SiliconFlow()
     for a in model.predict("你好", stream=True):

src/routers/chat_router.py

@@ -6,6 +6,7 @@ from concurrent.futures import ThreadPoolExecutor
 from src.core import HistoryManager
 from src.core.startup import startup
 from src.utils.logging_config import setup_logger
+from src.utils.web_search import WebSearcher
 
 chat = APIRouter(prefix="/chat")
 logger = setup_logger("server-chat")
@@ -13,6 +14,7 @@ logger = setup_logger("server-chat")
 executor = ThreadPoolExecutor()
 
 refs_pool = {}
+web_searcher = WebSearcher()
 
 @chat.get("/")
 async def chat_get():
@@ -37,19 +39,45 @@ def chat_post(
         }, ensure_ascii=False).encode('utf-8') + b"\n"
 
     def generate_response():
+        modified_query = query
         
-        if meta.get("enable_retrieval"):
+        # 处理网页搜索
+        if meta and meta.get("enable_web_search"):
+            chunk = make_chunk("正在进行网络搜索...", "searching", history=None)
+            yield chunk
+
+            try:
+                search_results = web_searcher.search(query)
+                if search_results:
+                    search_context = web_searcher.format_search_results(search_results)
+                    # 将搜索结果添加到查询中
+                    modified_query = f"""基于以下网络搜索结果回答问题：
+
+                    {search_context}
+
+                    用户问题：{query}
+
+                    请综合以上搜索结果，给出准确、客观的回答。如果搜索结果与问题相关性不大，请直接基于你的知识回答。
+                    """
+                    logger.info(f"Web search results added to query")
+                else:
+                    logger.warning("No web search results found")
+            except Exception as e:
+                logger.error(f"Web search error: {str(e)}")
+                chunk = make_chunk("网络搜索失败，将直接回答问题。", "loading", history=None)
+                yield chunk
+
+        # 处理知识库检索
+        if meta and meta.get("enable_retrieval"):
             chunk = make_chunk("", "searching", history=None)
             yield chunk
 
-            new_query, refs = startup.retriever(query, history_manager.messages, meta)
+            modified_query, refs = startup.retriever(modified_query, history_manager.messages, meta)
             refs_pool[cur_res_id] = refs
-        else:
-            new_query = query
 
-        messages = history_manager.get_history_with_msg(new_query, max_rounds=meta.get('history_round'))
+        messages = history_manager.get_history_with_msg(modified_query, max_rounds=meta.get('history_round'))
-        history_manager.add_user(query)
+        history_manager.add_user(query)  # 注意这里使用原始查询
-        logger.debug(f"Web history: {history_manager.messages}")
+        logger.debug(f"Final query: {modified_query}")
 
         content = ""
         for delta in startup.model.predict(messages, stream=True):

web/src/components/ChatComponent.vue

@@ -81,6 +81,9 @@
           <div class="flex-center" @click="meta.use_web = !meta.use_web" v-if="configStore.config.enable_search_engine && meta.enable_retrieval">
             搜索引擎（Bing） <div @click.stop><a-switch v-model:checked="meta.use_web" /></div>
           </div>
+          <div class="flex-center" @click="meta.enable_web_search = !meta.enable_web_search">
+            网页搜索 <div @click.stop><a-switch v-model:checked="meta.enable_web_search" /></div>
+          </div>
           <div class="flex-center" v-if="configStore.config.enable_knowledge_base && meta.enable_retrieval">
             重写查询 <a-segmented v-model:value="meta.rewriteQuery" :options="['off', 'on', 'hyde']"/>
           </div>
@@ -166,6 +169,8 @@ import {
   GlobalOutlined,
   FileTextOutlined,
   RobotOutlined,
+  EditOutlined,
+  PlusOutlined,
 } from '@ant-design/icons-vue'
 import { onClickOutside } from '@vueuse/core'
 import { Marked } from 'marked';
@@ -210,12 +215,14 @@ const meta = reactive(JSON.parse(localStorage.getItem('meta')) || {
   enable_retrieval: false,
   use_graph: false,
   use_web: false,
+  enable_web_search: false,
   graph_name: "neo4j",
   rewriteQuery: "off",
   selectedKB: null,
   stream: true,
   summary_title: true,
   history_round: 5,
+  db_name: null,
 })
 
 const marked = new Marked(
@@ -327,35 +334,26 @@ const appendAiMessage = (message, refs=null) => {
 
 const updateMessage = (info) => {
   const message = conv.value.messages.find((message) => message.id === info.id);
-
   if (message) {
-    // 只有在 text 不为空时更新
+    try {
       if (info.text !== null && info.text !== undefined && info.text !== '') {
         message.text = info.text;
       }
-
-    // 只有在 refs 不为空时更新
-    if (info.refs !== null && info.refs !== undefined) {
-      message.refs = info.refs;
-    }
-
-    if (info.model_name !== null && info.model_name !== undefined && info.model_name !== '') {
-      message.model_name = info.model_name;
-    }
-
-    // 只有在 status 不为空时更新
       if (info.status !== null && info.status !== undefined && info.status !== '') {
         message.status = info.status;
       }
-
       if (info.meta !== null && info.meta !== undefined) {
         message.meta = info.meta;
       }
-  } else {
-    console.error('Message not found');
-  }
-
       scrollToBottom();
+    } catch (error) {
+      console.error('Error updating message:', error);
+      message.status = 'error';
+      message.text = '消息更新失败';
+    }
+  } else {
+    console.error('Message not found:', info.id);
+  }
 };
 
 
@@ -399,21 +397,21 @@ const loadDatabases = () => {
 }
 
 // 新函数用于处理 fetch 请求
-const fetchChatResponse = (user_input, cur_res_id) => {
+const fetchChatResponse = (requestData) => {
   fetch('/api/chat/', {
     method: 'POST',
-    body: JSON.stringify({
-      query: user_input,
-      history: conv.value.history,
-      meta: meta,
-      cur_res_id: cur_res_id,
-    }),
     headers: {
       'Content-Type': 'application/json'
-    }
+    },
+    body: JSON.stringify(requestData)
   })
-  .then((response) => {
+  .then(response => {
-    if (!response.body) throw new Error("ReadableStream not supported.");
+    if (!response.ok) {
+      throw new Error(`HTTP error! status: ${response.status}`);
+    }
+    if (!response.body) {
+      throw new Error("ReadableStream not supported.");
+    }
     const reader = response.body.getReader();
     const decoder = new TextDecoder("utf-8");
     let buffer = '';
@@ -421,22 +419,22 @@ const fetchChatResponse = (user_input, cur_res_id) => {
     const readChunk = () => {
       return reader.read().then(({ done, value }) => {
         if (done) {
-          const message = conv.value.messages.find((message) => message.id === cur_res_id)
+          const message = conv.value.messages.find((message) => message.id === requestData.cur_res_id)
           console.log(message)
           if (message.meta.enable_retrieval) {
             console.log("fetching refs")
-            fetchRefs(cur_res_id).then((data) => {
+            fetchRefs(requestData.cur_res_id).then((data) => {
               console.log(data)
               updateMessage({
-                id: cur_res_id,
+                id: requestData.cur_res_id,
                 refs: data,
                 status: "finished",
               });
-              groupRefs(cur_res_id);
+              groupRefs(requestData.cur_res_id);
             })
           } else {
             updateMessage({
-              id: cur_res_id,
+              id: requestData.cur_res_id,
               status: "finished",
             });
           }
@@ -455,7 +453,7 @@ const fetchChatResponse = (user_input, cur_res_id) => {
             try {
               const data = JSON.parse(line);
               updateMessage({
-                id: cur_res_id,
+                id: requestData.cur_res_id,
                 text: data.response,
                 model_name: data.model_name,
                 status: data.status,
@@ -482,12 +480,14 @@ const fetchChatResponse = (user_input, cur_res_id) => {
     readChunk();
   })
   .catch((error) => {
-    console.error(error);
+    console.error('Error in fetchChatResponse:', error);
     updateMessage({
-      id: cur_res_id,
+      id: requestData.cur_res_id,
       status: "error",
+      text: `请求错误：${error.message}`,
     });
     isStreaming.value = false;
+    message.error(`请求失败：${error.message}`);
   });
 }
 
@@ -518,9 +518,30 @@ const sendMessage = () => {
     appendAiMessage("", null);
     const cur_res_id = conv.value.messages[conv.value.messages.length - 1].id;
     conv.value.inputText = '';
-    meta.db_name = dbName;
     
-    fetchChatResponse(user_input, cur_res_id)
+    // 准备发送的数据
+    const requestData = {
+      query: user_input,
+      history: conv.value.history,
+      cur_res_id: cur_res_id,
+      meta: {
+        enable_retrieval: meta.enable_retrieval,
+        use_graph: meta.use_graph,
+        use_web: meta.use_web,
+        enable_web_search: meta.enable_web_search,
+        graph_name: meta.graph_name,
+        rewriteQuery: meta.rewriteQuery,
+        selectedKB: meta.selectedKB,
+        stream: meta.stream,
+        summary_title: meta.summary_title,
+        history_round: meta.history_round,
+        db_name: dbName,
+      }
+    };
+
+    console.log('Sending request with data:', requestData); // 添加日志
+
+    fetchChatResponse(requestData);
   } else {
     console.log('请输入消息');
   }
@@ -648,6 +669,17 @@ watch(
     &:hover {
       background-color: var(--main-light-3);
     }
+
+    .anticon {
+      margin-right: 8px;
+      font-size: 16px;
+    }
+    
+    .ant-switch {
+      &.ant-switch-checked {
+        background-color: var(--main-500);
+      }
+    }
   }
 }
 
@@ -977,7 +1009,15 @@ watch(
   }
 }
 
+.controls {
+  display: flex;
+  align-items: center;
+  gap: 8px;
   
+  .search-switch {
+    margin-right: 8px;
+  }
+}
 </style>
 
 <style lang="less">