diff --git a/.gitignore b/.gitignore index df19050c..439e3613 100644 --- a/.gitignore +++ b/.gitignore @@ -35,4 +35,7 @@ web/package-lock.json saves notebooks graphrag -docker/volumes \ No newline at end of file +docker/volumes + + +.cursorrules diff --git a/README.md b/README.md index 17daf2ac..49cf70c6 100644 --- a/README.md +++ b/README.md @@ -12,9 +12,6 @@ > [!NOTE] > 当前项目还处于开发的早期，还存在一些 BUG，有问题随时提 issue。 -已知问题： - -- [ ] 从 Flask 更换到 Fast API 之后，并行命令还存在问题。 ## 概述 diff --git a/requirements.txt b/requirements.txt index ed96727a..c7f87bf4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -24,4 +24,5 @@ opencv-python-headless docx2txt uvicorn[standard] fastapi -python-multipart \ No newline at end of file +python-multipart +tavily-python \ No newline at end of file diff --git a/src/config/__init__.py b/src/config/__init__.py index e19d86bc..236e6f3b 100644 --- a/src/config/__init__.py +++ b/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 diff --git a/src/config/base.yaml b/src/config/base.yaml new file mode 100644 index 00000000..54c3aa4d --- /dev/null +++ b/src/config/base.yaml @@ -0,0 +1,17 @@ +# 基础配置 +stream: true +save_dir: saves + +# 功能开关 +enable_reranker: false +enable_knowledge_base: false +enable_knowledge_graph: false +enable_search_engine: false +enable_web_search: false + +# 模型配置 +model_provider: "deepseek" # 设置为 deepseek +model_name: "deepseek-chat" # 设置默认模型名称 +embed_model: "zhipu-embedding-3" +reranker: "bge-reranker-v2-m3" +model_local_paths: {} \ No newline at end of file diff --git a/src/config/models.yaml b/src/config/models.yaml index 813f4f87..90be011f 100644 --- a/src/config/models.yaml +++ b/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 diff --git a/src/core/startup.py b/src/core/startup.py index 943dd10b..ee4f95dd 100644 --- a/src/core/startup.py +++ b/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) diff --git a/src/models/__init__.py b/src/models/__init__.py index b5efa9d0..4f7a3d39 100644 --- a/src/models/__init__.py +++ b/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 - - 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}") diff --git a/src/models/chat_model.py b/src/models/chat_model.py index 8f1924e7..bcf331b3 100644 --- a/src/models/chat_model.py +++ b/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") - 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): diff --git a/src/models/ollama_embedding.py b/src/models/ollama_embedding.py new file mode 100644 index 00000000..1b099be9 --- /dev/null +++ b/src/models/ollama_embedding.py @@ -0,0 +1,179 @@ +import os +import requests +import numpy as np +from typing import List, Union, Dict +from src.utils.logging_config import setup_logger + +logger = setup_logger("OllamaEmbedding") + +class OllamaEmbedding: + """ + 使用 Ollama API 进行文本嵌入的类 + """ + def __init__(self, model_info: Dict, config) -> None: + """ + 初始化 Ollama Embedding 模型 + + Args: + model_info: 模型信息字典 + config: 配置对象 + """ + self.config = config + self.model_info = model_info + self.base_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434") + self.model_name = model_info.get("name", "nomic-embed-text") + self.query_instruction_for_retrieval = "为这个句子生成表示以用于检索相关文章：" + logger.info(f"Ollama Embedding model {self.model_name} initialized") + + def _get_embedding(self, text: str) -> List[float]: + """ + 获取单个文本的嵌入向量 + + Args: + text: 输入文本 + + Returns: + 嵌入向量 + """ + url = f"{self.base_url}/api/embeddings" + try: + response = requests.post(url, json={ + "model": self.model_name, + "prompt": text + }) + response.raise_for_status() + return response.json()["embedding"] + except Exception as e: + logger.error(f"Error getting embedding: {str(e)}") + raise + + def predict(self, messages: List[str]) -> List[List[float]]: + """ + 批量获取文本嵌入向量 + + Args: + messages: 文本列表 + + Returns: + 嵌入向量列表 + """ + embeddings = [] + batch_size = 20 + + for i in range(0, len(messages), batch_size): + batch = messages[i:i + batch_size] + logger.info(f"Processing batch {i//batch_size + 1}, size: {len(batch)}") + + batch_embeddings = [] + for text in batch: + embedding = self._get_embedding(text) + batch_embeddings.append(embedding) + + embeddings.extend(batch_embeddings) + + return embeddings + + def encode(self, messages: Union[str, List[str]]) -> List[List[float]]: + """ + 编码文本 + + Args: + messages: 单个文本或文本列表 + + Returns: + 嵌入向量列表 + """ + if isinstance(messages, str): + messages = [messages] + return self.predict(messages) + + def encode_queries(self, queries: List[str]) -> List[List[float]]: + """ + 编码查询文本 + + Args: + queries: 查询文本列表 + + Returns: + 查询文本的嵌入向量列表 + """ + return self.predict(queries) + + +class OllamaReranker: + """ + 使用 Ollama API 进行文本重排序的类 + """ + def __init__(self, config) -> None: + """ + 初始化 Ollama Reranker + + Args: + config: 配置对象 + """ + self.config = config + self.base_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434") + self.model_name = config.reranker + logger.info(f"Ollama Reranker model {self.model_name} initialized") + + def compute_score(self, query: str, passage: str) -> float: + """ + 计算查询和文本段落之间的相关性分数 + + Args: + query: 查询文本 + passage: 段落文本 + + Returns: + 相关性分数 + """ + prompt = f"Query: {query}\nPassage: {passage}\nRate the relevance of the passage to the query on a scale of 0 to 1:" + + try: + response = requests.post( + f"{self.base_url}/api/generate", + json={ + "model": self.model_name, + "prompt": prompt, + "stream": False + } + ) + response.raise_for_status() + + # 提取生成的数字作为分数 + result = response.json()["response"].strip() + try: + score = float(result) + return min(max(score, 0), 1) # 确保分数在 0-1 之间 + except ValueError: + logger.warning(f"Could not parse score from response: {result}") + return 0.0 + + except Exception as e: + logger.error(f"Error computing rerank score: {str(e)}") + return 0.0 + + def rerank(self, query: str, passages: List[str], top_n: int = None) -> List[Dict]: + """ + 重新排序文本段落 + + Args: + query: 查询文本 + passages: 段落文本列表 + top_n: 返回前 n 个结果 + + Returns: + 排序后的结果列表，每个元素包含索引和分数 + """ + scores = [] + for i, passage in enumerate(passages): + score = self.compute_score(query, passage) + scores.append({"index": i, "score": score}) + + # 按分数降序排序 + sorted_results = sorted(scores, key=lambda x: x["score"], reverse=True) + + if top_n: + sorted_results = sorted_results[:top_n] + + return sorted_results \ No newline at end of file diff --git a/src/routers/chat_router.py b/src/routers/chat_router.py index 1ef7ed5e..3f7d6117 100644 --- a/src/routers/chat_router.py +++ b/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,20 +39,50 @@ 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 +<<<<<<< HEAD meta["config"] = startup.config new_query, refs = startup.retriever(query, history_manager.messages, meta) +======= + modified_query, refs = startup.retriever(modified_query, history_manager.messages, meta) +>>>>>>> feature/web_search 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): diff --git a/src/utils/web_search.py b/src/utils/web_search.py new file mode 100644 index 00000000..7c5062b4 --- /dev/null +++ b/src/utils/web_search.py @@ -0,0 +1,69 @@ +import os +from typing import List, Dict +from tavily import TavilyClient +from src.utils.logging_config import setup_logger + +logger = setup_logger("web-search") + +class WebSearcher: + def __init__(self): + api_key = os.getenv("TAVILY_API_KEY", "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") + if not api_key: + raise ValueError("TAVILY_API_KEY environment variable is not set") + self.client = TavilyClient(api_key) + logger.info("WebSearcher initialized with Tavily client") + + def search(self, query: str, max_results: int = 1) -> List[Dict]: + """ + 使用 Tavily 搜索相关内容 + + Args: + query: 搜索查询 + max_results: 最大返回结果数 + + Returns: + 搜索结果列表 + """ + try: + search_results = self.client.search( + query=query, + search_depth="basic", + max_results=max_results + ) + + # 提取需要的信息 + formatted_results = [] + for result in search_results['results'][:max_results]: + formatted_results.append({ + 'title': result.get('title', ''), + 'content': result.get('content', ''), + 'url': result.get('url', ''), + 'score': result.get('score', 0) + }) + + return formatted_results + + except Exception as e: + logger.error(f"Error during web search: {str(e)}") + return [] + + def format_search_results(self, results: List[Dict]) -> str: + """ + 将搜索结果格式化为文本 + + Args: + results: 搜索结果列表 + + Returns: + 格式化后的文本 + """ + if not results: + return "没有找到相关的网络搜索结果。" + + formatted_text = "以下是相关的网络搜索结果：\n\n" + for i, result in enumerate(results, 1): + formatted_text += f"{i}. {result['title']}\n" + formatted_text += f" {result['content']}\n" + formatted_text += f" 来源: {result['url']}\n\n" + + return formatted_text \ No newline at end of file diff --git a/web/src/components/ChatComponent.vue b/web/src/components/ChatComponent.vue index d9ea6fe8..d34d709e 100644 --- a/web/src/components/ChatComponent.vue +++ b/web/src/components/ChatComponent.vue @@ -81,6 +81,9 @@
搜索引擎（Bing）
+
+ 网页搜索
+
@@ -166,6 +169,8 @@ import { GlobalOutlined, FileTextOutlined, RobotOutlined, + EditOutlined, + PlusOutlined, } from '@ant-design/icons-vue' import { onClickOutside } from '@vueuse/core' import { Marked } from 'marked'; @@ -206,12 +211,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", // use_rewrite_query: "off", selectedKB: null, stream: true, summary_title: true, history_round: 5, + db_name: null, }) const marked = new Marked( @@ -323,35 +330,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(); }; @@ -395,21 +393,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 = ''; @@ -417,22 +415,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", }); } @@ -451,7 +449,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, @@ -478,12 +476,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}`); }); } @@ -514,9 +514,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('请输入消息'); } @@ -644,6 +665,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); + } + } } } @@ -973,7 +1005,15 @@ watch( } } - +.controls { + display: flex; + align-items: center; + gap: 8px; + + .search-switch { + margin-right: 8px; + } +}