kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
littlewwwhite 2025-02-15 19:29:32 +08:00
parent 73ff8f5064
commit 5453452452
9 changed files with 179 additions and 100 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
.gitignore vendored
View File

@ -35,4 +35,7 @@ web/package-lock.json
saves
notebooks
graphrag
docker/volumes
docker/volumes
.cursorrules

3
requirements.txt
View File

@ -24,4 +24,5 @@ opencv-python-headless
docx2txt
uvicorn[standard]
fastapi
python-multipart
python-multipart
tavily-python

2
src/config/__init__.py
View File

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

1
src/config/models.yaml
View File

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

22
src/core/startup.py
View File

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

59
src/models/__init__.py
View File

@ -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
logger.info(f"Selecting model from {model_provider} with {model_name}")
if model_provider == "deepseek":
from src.models.chat_model import DeepSeek
return DeepSeek(model_name)
elif model_provider == "zhipu":
from src.models.chat_model import Zhipu
return Zhipu(model_name)
elif model_provider == "qianfan":
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`")
"""
根据配置选择模型
"""
if config.model_provider == "deepseek":
return DeepSeek(config.model_name)
elif config.model_provider == "zhipu":
return Zhipu(config.model_name)
elif config.model_provider == "openai":
return OpenModel(config.model_name)
elif config.model_provider == "qianfan":
return Qianfan(config.model_name)
elif config.model_provider == "dashscope":
return DashScope(config.model_name)
elif config.model_provider == "siliconflow":
return SiliconFlow(config.model_name)
else:
raise ValueError(f"Model provider {model_provider} not supported")
raise ValueError(f"Unsupported model provider: {config.model_provider}")

13
src/models/chat_model.py
View File

@ -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")
base_url = "https://api.deepseek.com"
api_key = os.getenv("DEEPSEEK_API_KEY", "your-default-api-key")
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):

42
src/routers/chat_router.py
View File

@ -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 and meta.get("enable_web_search"):
chunk = make_chunk("正在进行网络搜索...", "searching", history=None)
yield chunk
if meta.get("enable_retrieval"):
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'))
history_manager.add_user(query)
logger.debug(f"Web history: {history_manager.messages}")
messages = history_manager.get_history_with_msg(modified_query, max_rounds=meta.get('history_round'))
history_manager.add_user(query) # 注意这里使用原始查询
logger.debug(f"Final query: {modified_query}")
content = ""
for delta in startup.model.predict(messages, stream=True):

132
web/src/components/ChatComponent.vue
View File

@ -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 不为空时更新
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;
try {
if (info.text !== null && info.text !== undefined && info.text !== '') {
message.text = info.text;
}
if (info.status !== null && info.status !== undefined && info.status !== '') {
message.status = info.status;
}
if (info.meta !== null && info.meta !== undefined) {
message.meta = info.meta;
}
scrollToBottom();
} catch (error) {
console.error('Error updating message:', error);
message.status = 'error';
message.text = '消息更新失败';
}
} else {
console.error('Message not found');
console.error('Message not found:', info.id);
}
scrollToBottom();
};
@ -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) => {
if (!response.body) throw new Error("ReadableStream not supported.");
.then(response => {
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;
// 准备发送的数据
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,
}
};
fetchChatResponse(user_input, cur_res_id)
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">