From 6810f1f84ead518229e6e5241b85e2133be51f8d Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Wed, 2 Jul 2025 02:38:36 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=9C=A8=E4=BE=9D=E8=B5=96=E7=BB=84?= =?UTF-8?q?=E4=B8=AD=E6=B7=BB=E5=8A=A0=20ruff=20=E4=BE=9D=E8=B5=96?= =?UTF-8?q?=EF=BC=8C=E5=B9=B6=E4=BC=98=E5=8C=96=E4=BB=A3=E7=A0=81=E4=B8=AD?= =?UTF-8?q?=E7=9A=84=E7=B1=BB=E5=9E=8B=E6=8F=90=E7=A4=BA=E5=92=8C=E5=BC=82?= =?UTF-8?q?=E5=B8=B8=E5=A4=84=E7=90=86=EF=BC=8C=E6=8F=90=E5=8D=87=E4=BB=A3?= =?UTF-8?q?=E7=A0=81=E5=8F=AF=E8=AF=BB=E6=80=A7=E5=92=8C=E7=A8=B3=E5=AE=9A?= =?UTF-8?q?=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 1 + server/routers/graph_router.py | 50 +++--- src/agents/chatbot/configuration.py | 6 +- src/agents/utils.py | 62 ++++++-- src/config/__init__.py | 30 ---- src/core/graphbase.py | 1 - src/core/lightrag_based_kb.py | 20 +-- src/models/__init__.py | 41 +---- src/models/chat_model.py | 5 +- src/models/embedding.py | 126 +++------------ src/models/rerank_model.py | 48 +++--- src/static/models.yaml | 17 ++- .../components/ModelProvidersComponent.vue | 3 +- web/src/views/DataBaseInfoView.vue | 143 +++--------------- 14 files changed, 164 insertions(+), 389 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a15cac71..3830fe67 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ lint.ignore = ["F401"] # 忽略的规则 [dependency-groups] dev = [ + "ruff>=0.12.1", "vllm>=0.8.5.post1", ] diff --git a/server/routers/graph_router.py b/server/routers/graph_router.py index 88972f42..00eb546d 100644 --- a/server/routers/graph_router.py +++ b/server/routers/graph_router.py @@ -35,19 +35,19 @@ async def get_subgraph( """ try: logger.info(f"获取子图数据 - db_id: {db_id}, node_label: {node_label}, max_depth: {max_depth}, max_nodes: {max_nodes}") - + # 获取 LightRAG 实例 rag_instance = await knowledge_base._get_lightrag_instance(db_id) if not rag_instance: raise HTTPException(status_code=404, detail=f"数据库 {db_id} 不存在") - + # 使用 LightRAG 的原生 get_knowledge_graph 方法 knowledge_graph = await rag_instance.get_knowledge_graph( node_label=node_label, max_depth=max_depth, max_nodes=max_nodes ) - + # 将 LightRAG 的 KnowledgeGraph 格式转换为前端需要的格式 nodes = [] for node in knowledge_graph.nodes: @@ -57,7 +57,7 @@ async def get_subgraph( "entity_type": node.properties.get("entity_type", "unknown"), "properties": node.properties }) - + edges = [] for edge in knowledge_graph.edges: edges.append({ @@ -67,7 +67,7 @@ async def get_subgraph( "type": edge.type, "properties": edge.properties }) - + result = { "success": True, "data": { @@ -78,10 +78,10 @@ async def get_subgraph( "total_edges": len(edges) } } - + logger.info(f"成功获取子图 - 节点数: {len(nodes)}, 边数: {len(edges)}") return result - + except Exception as e: logger.error(f"获取子图数据失败: {e}") logger.error(f"Traceback: {traceback.format_exc()}") @@ -104,22 +104,22 @@ async def get_graph_labels( """ try: logger.info(f"获取图谱标签 - db_id: {db_id}") - + # 获取 LightRAG 实例 rag_instance = await knowledge_base._get_lightrag_instance(db_id) if not rag_instance: raise HTTPException(status_code=404, detail=f"数据库 {db_id} 不存在") - + # 使用 LightRAG 的原生方法获取所有标签 labels = await rag_instance.get_graph_labels() - + return { "success": True, "data": { "labels": labels } } - + except Exception as e: logger.error(f"获取图谱标签失败: {e}") logger.error(f"Traceback: {traceback.format_exc()}") @@ -142,7 +142,7 @@ async def get_available_databases( "success": True, "data": databases } - + except Exception as e: logger.error(f"获取数据库列表失败: {e}") raise HTTPException(status_code=500, detail=f"获取数据库列表失败: {str(e)}") @@ -154,8 +154,8 @@ async def get_graph_nodes_legacy( db_id: str = Query(..., description="数据库ID"), limit: int = Query(500, description="最大节点数量", ge=1, le=2000), offset: int = Query(0, description="偏移量", ge=0), - entity_type: Optional[str] = Query(None, description="实体类型筛选"), - search: Optional[str] = Query(None, description="搜索关键词"), + entity_type: str | None = Query(None, description="实体类型筛选"), + search: str | None = Query(None, description="搜索关键词"), current_user: User = Depends(get_admin_user) ): """ @@ -173,7 +173,7 @@ async def get_graph_nodes_legacy( "total": 0 } } - + except Exception as e: logger.error(f"获取图节点数据失败: {e}") raise HTTPException(status_code=500, detail=f"获取图节点数据失败: {str(e)}") @@ -184,7 +184,7 @@ async def get_graph_edges_legacy( db_id: str = Query(..., description="数据库ID"), limit: int = Query(500, description="最大边数量", ge=1, le=2000), offset: int = Query(0, description="偏移量", ge=0), - min_weight: Optional[float] = Query(None, description="最小权重筛选"), + min_weight: float | None = Query(None, description="最小权重筛选"), current_user: User = Depends(get_admin_user) ): """ @@ -202,7 +202,7 @@ async def get_graph_edges_legacy( "total": 0 } } - + except Exception as e: logger.error(f"获取图边数据失败: {e}") raise HTTPException(status_code=500, detail=f"获取图边数据失败: {str(e)}") @@ -218,30 +218,30 @@ async def get_graph_stats( """ try: logger.info(f"获取图谱统计信息 - db_id: {db_id}") - + # 获取 LightRAG 实例 rag_instance = await knowledge_base._get_lightrag_instance(db_id) if not rag_instance: raise HTTPException(status_code=404, detail=f"数据库 {db_id} 不存在") - + # 通过获取全图来统计节点和边的数量 knowledge_graph = await rag_instance.get_knowledge_graph( node_label="*", max_depth=1, max_nodes=10000 # 设置较大值以获取完整统计 ) - + # 统计实体类型分布 entity_types = {} for node in knowledge_graph.nodes: entity_type = node.properties.get("entity_type", "unknown") entity_types[entity_type] = entity_types.get(entity_type, 0) + 1 - + entity_types_list = [ - {"type": k, "count": v} + {"type": k, "count": v} for k, v in sorted(entity_types.items(), key=lambda x: x[1], reverse=True) ] - + return { "success": True, "data": { @@ -251,8 +251,8 @@ async def get_graph_stats( "is_truncated": knowledge_graph.is_truncated } } - + except Exception as e: logger.error(f"获取图谱统计信息失败: {e}") logger.error(f"Traceback: {traceback.format_exc()}") - raise HTTPException(status_code=500, detail=f"获取图谱统计信息失败: {str(e)}") \ No newline at end of file + raise HTTPException(status_code=500, detail=f"获取图谱统计信息失败: {str(e)}") diff --git a/src/agents/chatbot/configuration.py b/src/agents/chatbot/configuration.py index 57c5847e..2180311e 100644 --- a/src/agents/chatbot/configuration.py +++ b/src/agents/chatbot/configuration.py @@ -29,11 +29,7 @@ class ChatbotConfiguration(Configuration): metadata={ "name": "智能体模型", "configurable": True, - "options": [ - "zhipu/glm-4-plus", - "siliconflow/Qwen/Qwen2.5-72B-Instruct", - "siliconflow/Qwen/Qwen2.5-7B-Instruct", - ], + "options": [], "description": "智能体的驱动模型" }, ) diff --git a/src/agents/utils.py b/src/agents/utils.py index 9872fbf7..44b65c64 100644 --- a/src/agents/utils.py +++ b/src/agents/utils.py @@ -1,29 +1,73 @@ from datetime import datetime, timezone, UTC import asyncio +import os +import traceback -from src.models import select_model +from src import config +from src.utils import logger, get_docker_safe_url +from src.models import get_custom_model from src.agents.registry import BaseAgent from langchain_core.language_models import BaseChatModel from langchain_core.runnables import RunnableConfig from langchain_core.messages import AIMessageChunk, ToolMessage +from pydantic import SecretStr def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel: - """Load a chat model from a fully specified name. - - Args: - fully_specified_name (str): String in the format 'provider/model'. - **kwargs: Additional parameters to pass to the model. + """ + Load a chat model from a fully specified name. """ provider, model = fully_specified_name.split("/", maxsplit=1) - model_instance = select_model(model_name=model, model_provider=provider) - return model_instance.chat_open_ai + if provider == "custom": + from langchain_openai import ChatOpenAI + model_info = get_custom_model(model) + api_key = model_info.get("api_key") or "custom_model" + base_url = get_docker_safe_url(model_info["api_base"]) + model_name = model_info.get("name") or "custom_model" + return ChatOpenAI( + model=model_name, + api_key=SecretStr(api_key), + base_url=base_url, + ) + + model_info = config.model_names.get(provider, {}) + api_key = os.getenv(model_info["env"][0], model_info["env"][0]) + base_url = get_docker_safe_url(model_info["base_url"]) + + if provider in ["deepseek", "dashscope"]: + from langchain_deepseek import ChatDeepSeek + return ChatDeepSeek( + model=model, + api_key=SecretStr(api_key), + base_url=base_url, + api_base=base_url, + ) + + elif provider == "together": + from langchain_together import ChatTogether + return ChatTogether( + model=model, + api_key=SecretStr(api_key), + base_url=base_url, + ) + + else: + try: # 其他模型，默认使用OpenAIBase, like openai, zhipuai + from langchain_openai import ChatOpenAI + return ChatOpenAI( + model=model, + api_key=SecretStr(api_key), + base_url=base_url, + ) + except Exception as e: + raise ValueError(f"Model provider {provider} load failed, {e} \n {traceback.format_exc()}") -async def agent_cli(agent: BaseAgent, config: RunnableConfig = None): + +async def agent_cli(agent: BaseAgent, config: RunnableConfig | None = None): config = config or {} if "configurable" not in config: config["configurable"] = {} diff --git a/src/config/__init__.py b/src/config/__init__.py index b21f0732..89d7b35c 100644 --- a/src/config/__init__.py +++ b/src/config/__init__.py @@ -59,9 +59,6 @@ class Config(SimpleConfig): self.add_item("embed_model", default="siliconflow/BAAI/bge-m3", des="Embedding 模型", choices=list(self.embed_model_names.keys())) self.add_item("reranker", default="siliconflow/BAAI/bge-reranker-v2-m3", des="Re-Ranker 模型", choices=list(self.reranker_names.keys())) # noqa: E501 - self.add_item("model_local_paths", default={}, des="本地模型路径") - self.add_item("use_rewrite_query", default="on", des="重写查询", choices=["off", "on", "hyde"]) - self.add_item("device", default="cuda", des="运行本地模型的设备", choices=["cpu", "cuda"]) ### <<< 默认配置结束 self.load() @@ -120,7 +117,6 @@ class Config(SimpleConfig): """ 处理配置 """ - model_provider_info = self.model_names.get(self.model_provider, {}) self.model_dir = os.environ.get("MODEL_DIR", "") if self.model_dir: @@ -130,26 +126,6 @@ class Config(SimpleConfig): logger.warning(f"Warning: The model directory （{self.model_dir}） does not exist. If not configured, please ignore it. If configured, please check if the configuration is correct;" "For example, the mapping in the docker-compose file") - - # 检查模型提供商是否存在 - if self.model_provider != "custom": - if self.model_name not in model_provider_info["models"]: - logger.warning(f"Model name {self.model_name} not in {self.model_provider}, using default model name") - self.model_name = model_provider_info["default"] - - default_model_name = model_provider_info["default"] - self.model_name = self.get("model_name") or default_model_name - else: - self.model_name = self.get("model_name") - if self.model_name not in [item["custom_id"] for item in self.get("custom_models", [])]: - logger.warning(f"Model name {self.model_name} not in custom models, using default model name") - if self.get("custom_models", []): - self.model_name = self.get("custom_models", [])[0]["custom_id"] - else: - self.model_name = self._config_items["model_name"]["default"] - self.model_provider = self._config_items["model_provider"]["default"] - logger.error(f"No custom models found, using default model {self.model_name} from {self.model_provider}") - # 检查模型提供商的环境变量 conds = {} self.model_provider_status = {} @@ -158,12 +134,6 @@ class Config(SimpleConfig): conds_bool = [bool(os.getenv(_k)) for _k in conds[provider]] self.model_provider_status[provider] = all(conds_bool) - # 检查web_search的环境变量 - # if self.enable_web_search and not os.getenv("TAVILY_API_KEY"): - # logger.warning("TAVILY_API_KEY not set, web search will be disabled") - # self.enable_web_search = False - - # 2025.04.08 修改为不手动配置，只要配置了TAVILY_API_KEY，就默认开启web_search if os.getenv("TAVILY_API_KEY"): self.enable_web_search = True diff --git a/src/core/graphbase.py b/src/core/graphbase.py index f45dd446..21a0a1ae 100644 --- a/src/core/graphbase.py +++ b/src/core/graphbase.py @@ -3,7 +3,6 @@ import json import warnings import traceback -import torch from neo4j import GraphDatabase as GD from src import config diff --git a/src/core/lightrag_based_kb.py b/src/core/lightrag_based_kb.py index 9df55877..68b7c109 100644 --- a/src/core/lightrag_based_kb.py +++ b/src/core/lightrag_based_kb.py @@ -5,7 +5,7 @@ import traceback import shutil import asyncio from pathlib import Path -from typing import Dict, List, Any, Optional +from typing import Optional from datetime import datetime from lightrag import LightRAG, QueryParam @@ -26,11 +26,11 @@ class LightRagBasedKB: def __init__(self) -> None: # 存储 LightRAG 实例映射 {db_id: LightRAG} - self.instances: Dict[str, LightRAG] = {} + self.instances: dict[str, LightRAG] = {} # 数据库元信息存储 {db_id: metadata} - self.databases_meta: Dict[str, dict] = {} + self.databases_meta: dict[str, dict] = {} # 文件信息存储 {file_id: file_info} - self.files_meta: Dict[str, dict] = {} + self.files_meta: dict[str, dict] = {} # 工作目录 self.work_dir = os.path.join(config.save_dir, "lightrag_data") os.makedirs(self.work_dir, exist_ok=True) @@ -45,7 +45,7 @@ class LightRagBasedKB: meta_file = os.path.join(self.work_dir, "metadata.json") if os.path.exists(meta_file): try: - with open(meta_file, 'r', encoding='utf-8') as f: + with open(meta_file, encoding='utf-8') as f: data = json.load(f) self.databases_meta = data.get("databases", {}) self.files_meta = data.get("files", {}) @@ -69,7 +69,7 @@ class LightRagBasedKB: except Exception as e: logger.error(f"Failed to save metadata: {e}") - async def _get_lightrag_instance(self, db_id: str) -> Optional[LightRAG]: + async def _get_lightrag_instance(self, db_id: str) -> LightRAG | None: """获取或创建 LightRAG 实例""" logger.info(f"Getting or creating LightRAG instance for {db_id}") @@ -170,13 +170,13 @@ class LightRagBasedKB: elif file_ext in ['.txt', '.md']: # 直接读取文本文件 - with open(file_path_obj, 'r', encoding='utf-8') as f: + with open(file_path_obj, encoding='utf-8') as f: content = f.read() return f"# {file_path_obj.name}\n\n{content}" elif file_ext in ['.doc', '.docx']: # 处理 Word 文档 - + from docx import Document # type: ignore doc = Document(file_path_obj) text = '\n'.join([para.text for para in doc.paragraphs]) @@ -416,7 +416,7 @@ class LightRagBasedKB: try: # 获取文档的所有 chunks assert hasattr(rag.text_chunks, 'get_all'), "text_chunks does not have get_all method" - all_chunks = await rag.text_chunks.get_all() + all_chunks = await rag.text_chunks.get_all() # type: ignore # 筛选属于该文档的 chunks doc_chunks = [] @@ -506,4 +506,4 @@ class LightRagBasedKB: "description": meta["description"], "retriever": make_retriever(db_id), } - return retrievers \ No newline at end of file + return retrievers diff --git a/src/models/__init__.py b/src/models/__init__.py index 86b31bfe..885ebdb1 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -5,19 +5,14 @@ from src import config from src.utils.logging_config import logger from src.models.chat_model import OpenAIBase -def select_model(model_provider=None, model_name=None): +def select_model(model_provider, model_name=None): """根据模型提供者选择模型""" - model_provider = model_provider or config.model_provider + assert model_provider is not None, "Model provider not specified" model_info = config.model_names.get(model_provider, {}) - model_name = model_name or config.model_name or model_info.get("default", "") - + model_name = model_name or model_info.get("default", "") logger.info(f"Selecting model from `{model_provider}` with `{model_name}`") - - if model_provider is None: - raise ValueError("Model provider not specified, please modify `model_provider` in `src/config/base.yaml`") - if model_provider == "qianfan": from src.models.chat_model import Qianfan return Qianfan(model_name) @@ -26,33 +21,6 @@ def select_model(model_provider=None, model_name=None): from src.models.chat_model import OpenModel return OpenModel(model_name) - if model_provider == "deepseek" or model_provider == "dashscope": - from langchain_deepseek import ChatDeepSeek - return OpenAIBase( - api_key=os.getenv(model_info["env"][0]), - base_url=model_info["base_url"], - model_name=model_name, - chat_open_ai=ChatDeepSeek( - model=model_name, - api_key=os.getenv(model_info["env"][0]), - base_url=model_info["base_url"], - api_base=model_info["base_url"], - ) - ) - - if model_provider == "together": - from langchain_together import ChatTogether - return OpenAIBase( - api_key=os.getenv(model_info["env"][0]), - base_url=model_info["base_url"], - model_name=model_name, - chat_open_ai=ChatTogether( - model=model_name, - api_key=os.getenv(model_info["env"][0]), - base_url=model_info["base_url"], - ) - ) - if model_provider == "custom": model_info = get_custom_model(model_name) @@ -69,8 +37,7 @@ def select_model(model_provider=None, model_name=None): return model except Exception as e: raise ValueError(f"Model provider {model_provider} load failed, {e} \n {traceback.format_exc()}") - - + def get_custom_model(model_id): """return model_info""" assert config.custom_models is not None, "custom_models is not set" diff --git a/src/models/chat_model.py b/src/models/chat_model.py index d8e9f28a..e3a97c79 100644 --- a/src/models/chat_model.py +++ b/src/models/chat_model.py @@ -5,15 +5,12 @@ from src.utils import logger, get_docker_safe_url from langchain_openai import ChatOpenAI class OpenAIBase: - def __init__(self, api_key, base_url, model_name, chat_open_ai=None, **kwargs): + def __init__(self, api_key, base_url, model_name, **kwargs): self.api_key = api_key self.base_url = base_url self.client = OpenAI(api_key=api_key, base_url=base_url) self.model_name = model_name self.info = kwargs - self.chat_open_ai = chat_open_ai or ChatOpenAI(model=model_name, - api_key=api_key, - base_url=base_url) def predict(self, message, stream=False): if isinstance(message, str): diff --git a/src/models/embedding.py b/src/models/embedding.py index f4823cc7..9cfb9afb 100644 --- a/src/models/embedding.py +++ b/src/models/embedding.py @@ -13,19 +13,18 @@ from src.utils import hashstr, logger, get_docker_safe_url class BaseEmbeddingModel: embed_state = {} + def __init__(self, model_id): + self.model_id = model_id + self.info = config.embed_model_names[model_id] + self.model = self.info["name"] + self.dimension = self.info.get("dimension", None) + self.url = get_docker_safe_url(self.info["base_url"]) + self.api_key = os.getenv(self.info["api_key"], self.info["api_key"]) + @abstractmethod def predict(self, message): raise NotImplementedError("Subclasses must implement this method") - def get_dimension(self): - if hasattr(self, "dimension"): - return self.dimension - - if hasattr(self, "embed_model_fullname"): - return config.embed_model_names[self.embed_model_fullname].get("dimension", None) - - return config.embed_model_names[self.model].get("dimension", None) - def encode(self, message): return self.predict(message) @@ -66,77 +65,14 @@ class BaseEmbeddingModel: return data -class LocalEmbeddingModel(BaseEmbeddingModel): - def __init__(self, **kwargs): - info = config.embed_model_names[config.embed_model] - - self.model = config.model_local_paths.get(info["name"], info.get("local_path")) - self.model = self.model or info["name"] - self.dimension = info["dimension"] - self.embed_model_fullname = config.embed_model - - if os.getenv("MODEL_DIR"): - if os.path.exists(_path := os.path.join(os.getenv("MODEL_DIR"), self.model)): - self.model = _path - else: - logger.warning(f"Local model `{info['name']}` not found in `{self.model}`, using `{info['name']}`") - - logger.info(f"Loading local model `{info['name']}` from `{self.model}` with device `{config.device}`") - logger.debug("如果没配置任何路径的话，正常情况下会自动从 Huggingface 下载模型，如果遇到下载失败，可以尝试使用 HF_MIRROR 环境变量；" - f"如果还是不行，建议手动下载到某个文件夹，比如 {os.getenv('MODEL_DIR', '/models')}/BAAI/bge-m3 目录下；") - - self.model = HuggingFaceEmbeddings( - model_name=self.model, - model_kwargs={'device': config.device}, - encode_kwargs={ - 'normalize_embeddings': True, - 'prompt_name': info.get("query_instruction", None), - }, - ) - - logger.info(f"Embedding model {info['name']} loaded, {self.model=}") - - def predict(self, message): - return self.model.embed_documents(message) - - async def aencode(self, message): - return await self.model.aembed_documents(message) - - def encode_queries(self, queries): - logger.warning("Huggingface Model 不支持批量 encode queries，因此使用训练实现") - data = [] - for q in queries: - data.append(self.predict(q)) - - return data - - -class ZhipuEmbedding(BaseEmbeddingModel): - - def __init__(self) -> None: - self.config = config - self.model = config.embed_model_names[config.embed_model]["name"] - self.dimension = config.embed_model_names[config.embed_model]["dimension"] - self.client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY")) - self.embed_model_fullname = config.embed_model - - def predict(self, message): - response = self.client.embeddings.create( - model=self.model, - input=message, - ) - data = [a.embedding for a in response.data] - return data - - class OllamaEmbedding(BaseEmbeddingModel): - def __init__(self) -> None: - self.info = config.embed_model_names[config.embed_model] - self.model = self.info["name"] - self.url = self.info.get("base_url", "http://localhost:11434/api/embed") - self.url = get_docker_safe_url(self.url) - self.dimension = self.info.get("dimension", None) - self.embed_model_fullname = config.embed_model + """ + Ollama Embedding Model + """ + + def __init__(self, model_id) -> None: + super().__init__(model_id) + self.url = self.url or get_docker_safe_url("http://localhost:11434/api/embed") def predict(self, message: list[str] | str): if isinstance(message, str): @@ -154,14 +90,8 @@ class OllamaEmbedding(BaseEmbeddingModel): class OtherEmbedding(BaseEmbeddingModel): - def __init__(self) -> None: - self.info = config.embed_model_names[config.embed_model] - self.embed_model_fullname = config.embed_model - self.dimension = self.info.get("dimension", None) - self.model = self.info["name"] - self.api_key = os.getenv(self.info["api_key"], self.info["api_key"]) - self.url = get_docker_safe_url(self.info["base_url"]) - assert self.url and self.model, f"URL and model are required. Cur embed model: {config.embed_model}" + def __init__(self, model_id) -> None: + super().__init__(model_id) self.headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json" @@ -181,26 +111,18 @@ class OtherEmbedding(BaseEmbeddingModel): "input": message, } -def get_embedding_model(): - provider, model_name = config.embed_model.split('/', 1) +def get_embedding_model(model_id): + provider, model_name = model_id.split('/', 1) if model_id else ("", "") support_embed_models = config.embed_model_names.keys() - assert config.embed_model in support_embed_models, f"Unsupported embed model: {config.embed_model}, only support {support_embed_models}" - logger.debug(f"Loading embedding model {config.embed_model}") + assert model_id in support_embed_models, f"Unsupported embed model: {model_id}, only support {support_embed_models}" + logger.debug(f"Loading embedding model {model_id}") if provider == "local": - logger.warning("[DEPRECATED] Local embedding model will be removed in v0.2, please use other embedding models") - model = LocalEmbeddingModel() - - elif provider == "zhipu": - model = ZhipuEmbedding() + raise ValueError("Local embedding model is not supported, please use other embedding models") elif provider == "ollama": - model = OllamaEmbedding() + model = OllamaEmbedding(model_id) else: - model = OtherEmbedding() + model = OtherEmbedding(model_id) return model - -def handle_local_model(paths, model_name, default_path): - model_path = paths.get(model_name, default_path) - return model_path diff --git a/src/models/rerank_model.py b/src/models/rerank_model.py index 01c4431d..f51944d7 100644 --- a/src/models/rerank_model.py +++ b/src/models/rerank_model.py @@ -2,34 +2,18 @@ import os import json import requests import numpy as np -from FlagEmbedding import FlagReranker from src import config from src.utils import logger, get_docker_safe_url - -class LocalReranker(FlagReranker): - def __init__(self, **kwargs): - model_info = config.reranker_names[config.reranker] - model_name_or_path = config.model_local_paths.get(model_info["name"], model_info.get("local_path")) - model_name_or_path = model_name_or_path or model_info["name"] - logger.info(f"Loading Reranker model {config.reranker} from {model_name_or_path}") - - super().__init__(model_name_or_path, use_fp16=True, device=config.device, **kwargs) - logger.info(f"Reranker model {config.reranker} loaded") - - def sigmoid(x): return 1 / (1 + np.exp(-x)) -class OnlineRerank: - def __init__(self, **kwargs): - model_info = config.reranker_names[config.reranker] - self.url = get_docker_safe_url(model_info["base_url"]) - self.model = model_info["name"] - - api_key = os.getenv(model_info["api_key"], model_info["api_key"]) - assert api_key, f"{model_info['name']} api_key is required" +class OnlineReranker: + def __init__(self, model_name, api_key, base_url, **kwargs): + self.url = get_docker_safe_url(base_url) + self.model = model_name + self.api_key = api_key self.headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json" @@ -59,15 +43,17 @@ class OnlineRerank: "max_chunks_per_doc": max_length, } -def get_reranker(): +def get_reranker(model_id, **kwargs): support_rerankers = config.reranker_names.keys() - assert config.reranker in support_rerankers, f"Unsupported Reranker: {config.reranker}, only support {support_rerankers}" - provider, model_name = config.reranker.split('/', 1) - if provider == "local": - logger.warning("[DEPRECATED] Local reranker will be removed in v0.2, please use other reranker") - return LocalReranker() - elif provider == "siliconflow": - return OnlineRerank() - else: - raise ValueError(f"Unsupported Reranker: {config.reranker}, only support {config.reranker_names.keys()}") + assert model_id in support_rerankers, f"Unsupported Reranker: {model_id}, only support {support_rerankers}" + model_info = config.reranker_names[model_id] + base_url = model_info["base_url"] + api_key = os.getenv(model_info["api_key"], model_info["api_key"]) + assert api_key, f"{model_info['name']} api_key is required" + return OnlineReranker( + model_name=model_info["name"], + api_key=api_key, + base_url=base_url, + **kwargs + ) diff --git a/src/static/models.yaml b/src/static/models.yaml index a1d9f4eb..e50fdd29 100644 --- a/src/static/models.yaml +++ b/src/static/models.yaml @@ -6,12 +6,11 @@ # ##################################################### - - MODEL_NAMES: openai: name: OpenAI url: https://platform.openai.com/docs/models + base_url: https://api.openai.com/v1 default: gpt-3.5-turbo env: - OPENAI_API_KEY @@ -20,21 +19,23 @@ MODEL_NAMES: - gpt-4o - gpt-4o-mini - gpt-3.5-turbo + deepseek: name: DeepSeek url: https://platform.deepseek.com/api-docs/zh-cn/pricing - default: deepseek-chat base_url: https://api.deepseek.com/v1 + default: deepseek-chat env: - DEEPSEEK_API_KEY models: - deepseek-chat - deepseek-reasoner + zhipu: name: 智谱AI (Zhipu) url: https://open.bigmodel.cn/dev/api - default: glm-4-flash base_url: https://open.bigmodel.cn/api/paas/v4/ + default: glm-4-flash env: - ZHIPUAI_API_KEY models: @@ -43,17 +44,19 @@ MODEL_NAMES: - glm-4-air - glm-4-flash - glm-z1-air + siliconflow: name: SiliconFlow url: https://cloud.siliconflow.cn/models - default: Qwen/Qwen2.5-7B-Instruct base_url: https://api.siliconflow.cn/v1 + default: Qwen/Qwen3-8B env: - SILICONFLOW_API_KEY models: - Pro/deepseek-ai/DeepSeek-R1 - Pro/deepseek-ai/DeepSeek-V3 - Qwen/QwQ-32B + - Qwen/Qwen3-8B together.ai: name: Together.ai @@ -71,8 +74,8 @@ MODEL_NAMES: dashscope: name: 阿里百炼 (DashScope) url: https://bailian.console.aliyun.com/?switchAgent=10226727&productCode=p_efm#/model-market - default: qwen3-235b-a22b base_url: https://dashscope.aliyuncs.com/compatible-mode/v1 + default: qwen3-235b-a22b env: - DASHSCOPE_API_KEY models: @@ -83,8 +86,8 @@ MODEL_NAMES: ark: name: 豆包（Ark） url: https://console.volcengine.com/ark/region:ark+cn-beijing/model - default: doubao-1-5-pro-32k-250115 base_url: https://ark.cn-beijing.volces.com/api/v3 + default: doubao-1-5-pro-32k-250115 env: - ARK_API_KEY models: diff --git a/web/src/components/ModelProvidersComponent.vue b/web/src/components/ModelProvidersComponent.vue index f503ce73..ce22939b 100644 --- a/web/src/components/ModelProvidersComponent.vue +++ b/web/src/components/ModelProvidersComponent.vue @@ -15,9 +15,8 @@ okText="确认删除" cancelText="取消" ok-type="danger" - :disabled="configStore.config.model_name == item.name" > - + diff --git a/web/src/views/DataBaseInfoView.vue b/web/src/views/DataBaseInfoView.vue index 6e068480..1865e706 100644 --- a/web/src/views/DataBaseInfoView.vue +++ b/web/src/views/DataBaseInfoView.vue @@ -15,7 +15,6 @@ - @@ -282,12 +281,12 @@
-

只使用上下文：

+

只使用提示：

筛选 TopK：

- +
@@ -349,7 +348,6 @@ import { CloseCircleFilled, ClockCircleFilled, DeleteOutlined, - CloudUploadOutlined, SearchOutlined, LoadingOutlined, FileOutlined, @@ -357,7 +355,6 @@ import { EditOutlined, PlusOutlined, SettingOutlined, - ApartmentOutlined, } from '@ant-design/icons-vue' import HeaderComponent from '@/components/HeaderComponent.vue'; import KnowledgeGraphViewer from '@/components/KnowledgeGraphViewer.vue'; @@ -376,7 +373,6 @@ const selectedFile = ref(null); // 查询测试 const queryText = ref(''); const queryResult = ref(null) -const filteredResults = ref([]) const configStore = useConfigStore() const state = reactive({ @@ -431,49 +427,7 @@ const pagination = ref({ onShowSizeChange: (current, pageSize) => pagination.value.pageSize = pageSize, }) -const filterQueryResults = () => { - if (!queryResult.value || !queryResult.value.all_results) { - return; - } - - let results = toRaw(queryResult.value.all_results); - console.log("results", results); - - if (meta.filter) { - results = results.filter(r => r.distance >= meta.distanceThreshold); - console.log("before", results); - - // 根据排序方式决定排序逻辑 - if (configStore.config.enable_reranker) { - // 先过滤掉低于阈值的结果 - results = results.filter(r => r.rerank_score >= meta.rerankThreshold); - - // 根据选择的排序方式进行排序 - if (meta.sortBy === 'rerank_score') { - results = results.sort((a, b) => b.rerank_score - a.rerank_score); - } else { - // 按距离排序 (数值越大表示越相似) - results = results.sort((a, b) => b.distance - a.distance); - } - } else { - // 没有启用重排序时，默认按距离排序 - results = results.sort((a, b) => b.distance - a.distance); - } - - console.log("after", results); - - results = results.slice(0, meta.topK); - } - - filteredResults.value = results; -} - const onQuery = () => { - // if (database.value.embed_model != configStore.config.embed_model) { - // message.error('向量模型不匹配，请重新选择') - // return - // } - console.log(queryText.value) state.searchLoading = true if (!queryText.value.trim()) { @@ -491,7 +445,6 @@ const onQuery = () => { .then(data => { console.log(data) queryResult.value = data - filterQueryResults() }) .catch(error => { console.error(error) @@ -665,7 +618,7 @@ const getDatabaseInfo = () => { const deleteFile = (fileId) => { state.lock = true console.debug("deleteFile", databaseId.value, fileId) - knowledgeBaseApi.deleteFile(databaseId.value, fileId) + return knowledgeBaseApi.deleteFile(databaseId.value, fileId) .then(data => { console.log(data) message.success(data.message || '删除成功') @@ -674,6 +627,7 @@ const deleteFile = (fileId) => { .catch(error => { console.error(error) message.error(error.message || '删除失败') + throw error }) .finally(() => { state.lock = false @@ -802,27 +756,7 @@ const addFiles = (items, contentType = 'file') => { }); }; -// "生成分块" - 修改后的逻辑（文件） -const handleFiles = () => { - console.log(fileList.value); - const files = fileList.value.filter(file => file.status === 'done').map(file => file.response.file_path); - console.log(files); - addFiles(files, 'file'); -} -// "生成分块" - 修改后的逻辑（URL） -const handleUrls = () => { - const urls = urlList.value.split('\n') - .map(url => url.trim()) - .filter(url => url.length > 0 && (url.startsWith('http://') || url.startsWith('https://'))); - - if (urls.length === 0) { - message.error('请输入有效的网页链接（必须以http://或https://开头）'); - return; - } - - addFiles(urls, 'url'); -}; const columns = [ // { title: '文件ID', dataIndex: 'file_id', key: 'file_id' }, @@ -843,10 +777,6 @@ watch(() => route.params.database_id, (newId) => { } ); -// 检测到 meta 变化时重新查询 -watch(() => meta, () => { - filterQueryResults() -}, { deep: true }) // 添加更多示例查询 const queryExamples = ref([ @@ -875,9 +805,20 @@ const urlList = ref(''); const chunkData = () => { if (uploadMode.value === 'file') { - handleFiles(); + const files = fileList.value.filter(file => file.status === 'done').map(file => file.response.file_path); + console.log(files); + addFiles(files, 'file'); } else if (uploadMode.value === 'url') { - handleUrls(); + const urls = urlList.value.split('\n') + .map(url => url.trim()) + .filter(url => url.length > 0 && (url.startsWith('http://') || url.startsWith('https://'))); + + if (urls.length === 0) { + message.error('请输入有效的网页链接（必须以http://或https://开头）'); + return; + } + + addFiles(urls, 'url'); } } @@ -1494,59 +1435,9 @@ const toggleAutoRefresh = (checked) => { } } -.chunk-preview { - margin-top: 20px; - .preview-header { - display: flex; - justify-content: space-between; - align-items: center; - margin-bottom: 16px; - h3 { - margin: 0; - color: var(--main-color); - font-size: 18px; - } - } - .result-cards { - display: grid; - grid-template-columns: repeat(auto-fill, minmax(600px, 1fr)); - gap: 12px; - margin-top: 10px; - } - - .chunk { - background-color: var(--main-light-5); - border: 1px solid var(--main-light-3); - border-radius: 8px; - padding: 16px; - word-wrap: break-word; - word-break: break-all; - transition: all 0.2s ease; - - &:hover { - background-color: var(--main-light-4); - border-color: var(--main-light-2); - box-shadow: 0 2px 6px rgba(0, 0, 0, 0.06); - } - - p { - margin: 0; - line-height: 1.6; - - strong { - color: var(--main-color); - margin-right: 6px; - } - } - } -} - -.url-input { - margin-bottom: 20px; -}