refactor: 重构配置/知识库模块，拆分为独立app.py文件
This commit is contained in:
Wenjie Zhang
parent
747c5de88f
commit
66c550c18c
1
.gitignore
vendored
1
.gitignore
vendored
@ -29,6 +29,7 @@ cache
|
||||
.vscode
|
||||
.idea
|
||||
.vibe
|
||||
.qoder
|
||||
*.nogit*
|
||||
*.private*
|
||||
*.local*
|
||||
|
||||
@ -3,14 +3,23 @@
|
||||
💭 **Features Todo**
|
||||
- [x] LangGraph 升级到 0.6+ 版本，并适配新特性，如 context 等，添加 MCP 工具的支持。
|
||||
- [x] 支持动态工具配置的同时，将 Configuration替换为 Context 后能够正常使用
|
||||
- [ ] 使用更好的知识图谱检索方法
|
||||
- [x] 使用更好的知识图谱检索方法（二跳结果查询）
|
||||
- [ ] 使用其他的聊天记录管理方法，解决两个问题，一个是上下文长度过长，一个是上下文的类型变得更加丰富，比如多模态等等。（现在是基于 LangGraph 的 Memory 实现的，v0.2.3 版本实现），暂定使用 [mem0](github.com/mem0ai/mem0) 来实现。但是目前了解下来，还不是我想要的那种方案。可能会基于这个实现一个 ThreadConvManager 这个类。
|
||||
- [ ] 添加对于上传文件的支持：这里的复杂的地方就在于如何和历史记录结合在一起（v0.2.3 版本实现，放在记忆管理后面）
|
||||
- [x] 将现在的 graphview.vue 文件中 GraphContainer相关的代码分离到一个单独的组件中，然后 actions 和 footer 都是作为 slot top/bottom 提供的
|
||||
- [x] 然后应用到 web/src/components/ToolCallingResult/KnowledgeGraphResult.vue 中，替换现有的 GraphContainer。
|
||||
- [x] 然后应用到 web/src/components/ToolCallingResult/KnowledgeGraphResult.vue 中，替换现有的 GraphContainer。
|
||||
- [ ] 知识图谱的上传和可视化，支持属性，标签的展示
|
||||
- [ ] 集成智能体评估，首先使用命令行来实现，然后考虑放在 UI 里面展示
|
||||
- [ ] 开发与生产环境隔离
|
||||
- [ ] 添加统计信息
|
||||
- [ ] 添加内容审查功能
|
||||
|
||||
📝 **Base**
|
||||
|
||||
- [ ] 优化全局配置的管理模型，以及配置文件的管理，子配置的管理
|
||||
- [ ] 新建 tasker 模块，用来管理所有的后台任务，UI 上使用侧边栏管理。
|
||||
- [ ] 新增 files 模块，用来管理文件上传，下载等
|
||||
|
||||
|
||||
🐛**BUGs**
|
||||
- [x] LlightRAG 知识库中，点击边，没有显示，但是在全屏的时候却又能够显示出来。
|
||||
@ -20,15 +29,17 @@
|
||||
- [x] 目前只能获取默认的 tools，单个智能体的tools没法展示
|
||||
- [x] 生成的过程中切换对话，会出现消息渲染混乱的问题，消息完全记载完成之后不会出现混乱
|
||||
- [ ] 当消息生成的时候有报错的时候，前端无显示
|
||||
- [ ] 偶然会出现 embedding 报错的情况，但不好复现
|
||||
- [x] GraphCanvas 在智能体消息页面有时候会出现渲染不出现的问题，可以添加一个slot 在右上角（加载 button），这样就可以强制重新渲染
|
||||
|
||||
💯 **More**:
|
||||
|
||||
下面的功能**可能**会放在后续版本实现，暂时未定
|
||||
|
||||
- [ ] 支持数据库查询功能
|
||||
- [x] 支持数据库查询功能
|
||||
- [x] 消息内容支持图片显示
|
||||
- [ ] 添加 SQL 读取工具
|
||||
- [ ] 添加绘图工具（这里的绘图是指绘制固定格式的图和表等，暂时没想好如何支持自定义绘图）
|
||||
- [x] 添加绘图工具（这里的绘图是指绘制固定格式的图和表等，暂时没想好如何支持自定义绘图）
|
||||
- [ ] 添加测试脚本，覆盖最常见的功能
|
||||
- [ ] 优化对文档信息的检索展示（检索结果页、详情页）
|
||||
- [ ] 集成 LangFuse (观望)添加用户日志与用户反馈模块，可以在 AgentView 中查看信息
|
||||
@ -5,30 +5,8 @@ from dotenv import load_dotenv
|
||||
load_dotenv("src/.env", override=True)
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor # noqa: E402
|
||||
from src.config import config # noqa: E402
|
||||
from src.knowledge import knowledge_base, graph_base # noqa: E402
|
||||
|
||||
|
||||
executor = ThreadPoolExecutor()
|
||||
|
||||
from src.config import Config # noqa: E402
|
||||
|
||||
config = Config()
|
||||
|
||||
# 导入知识库相关模块
|
||||
from src.knowledge.chroma_kb import ChromaKB # noqa: E402
|
||||
from src.knowledge.kb_factory import KnowledgeBaseFactory # noqa: E402
|
||||
from src.knowledge.kb_manager import KnowledgeBaseManager # noqa: E402
|
||||
from src.knowledge.lightrag_kb import LightRagKB # noqa: E402
|
||||
from src.knowledge.milvus_kb import MilvusKB # noqa: E402
|
||||
|
||||
# 注册知识库类型
|
||||
KnowledgeBaseFactory.register("chroma", ChromaKB, {"description": "基于 ChromaDB 的轻量级向量知识库，适合开发和小规模"})
|
||||
KnowledgeBaseFactory.register("milvus", MilvusKB, {"description": "基于 Milvus 的生产级向量知识库，适合高性能部署"})
|
||||
KnowledgeBaseFactory.register("lightrag", LightRagKB, {"description": "基于图检索的知识库，支持实体关系构建和复杂查询"})
|
||||
|
||||
|
||||
# 创建知识库管理器
|
||||
work_dir = os.path.join(config.save_dir, "knowledge_base_data")
|
||||
knowledge_base = KnowledgeBaseManager(work_dir)
|
||||
|
||||
from src.knowledge import GraphDatabase # noqa: E402
|
||||
|
||||
graph_base = GraphDatabase()
|
||||
|
||||
@ -1,228 +1,3 @@
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from .app import config
|
||||
|
||||
import yaml
|
||||
|
||||
from src.utils.logging_config import logger
|
||||
|
||||
DEFAULT_MOCK_API = "this_is_mock_api_key_in_frontend"
|
||||
|
||||
|
||||
class SimpleConfig(dict):
|
||||
def __key(self, key):
|
||||
return "" if key is None else key # 目前忘记了这里为什么要 lower 了，只能说配置项最好不要有大写的
|
||||
|
||||
def __str__(self):
|
||||
return json.dumps(self)
|
||||
|
||||
def __setattr__(self, key, value):
|
||||
self[self.__key(key)] = value
|
||||
|
||||
def __getattr__(self, key):
|
||||
return self.get(self.__key(key))
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self.get(self.__key(key))
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
return super().__setitem__(self.__key(key), value)
|
||||
|
||||
def __dict__(self):
|
||||
return {k: v for k, v in self.items()}
|
||||
|
||||
def update(self, other):
|
||||
for key, value in other.items():
|
||||
self[key] = value
|
||||
|
||||
|
||||
class Config(SimpleConfig):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._config_items = {}
|
||||
self.save_dir = os.getenv("SAVE_DIR", "saves")
|
||||
self.filename = str(Path(f"{self.save_dir}/config/base.yaml"))
|
||||
os.makedirs(os.path.dirname(self.filename), exist_ok=True)
|
||||
|
||||
self._update_models_from_file()
|
||||
|
||||
### >>> 默认配置
|
||||
# 功能选项
|
||||
self.add_item("enable_reranker", default=False, des="是否开启重排序")
|
||||
self.add_item(
|
||||
"enable_web_search",
|
||||
default=False,
|
||||
des=(
|
||||
"是否开启网页搜索（注：现阶段会根据 TAVILY_API_KEY 自动开启，无法手动配置，将会在下个版本移除此配置项）"
|
||||
),
|
||||
)
|
||||
# 默认智能体配置
|
||||
self.add_item("default_agent_id", default="", des="默认智能体ID")
|
||||
# 模型配置
|
||||
## 注意这里是模型名，而不是具体的模型路径，默认使用 HuggingFace 的路径
|
||||
## 如果需要自定义本地模型路径，则在 src/.env 中配置 MODEL_DIR
|
||||
self.add_item("model_provider", default="siliconflow", des="模型提供商", choices=list(self.model_names.keys()))
|
||||
self.add_item("model_name", default="zai-org/GLM-4.5", des="模型名称")
|
||||
|
||||
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.load()
|
||||
self.handle_self()
|
||||
|
||||
def add_item(self, key, default, des=None, choices=None):
|
||||
self.__setattr__(key, default)
|
||||
self._config_items[key] = {"default": default, "des": des, "choices": choices}
|
||||
|
||||
def __dict__(self):
|
||||
blocklist = [
|
||||
"_config_items",
|
||||
"model_names",
|
||||
"model_provider_status",
|
||||
"embed_model_names",
|
||||
"reranker_names",
|
||||
]
|
||||
return {k: v for k, v in self.items() if k not in blocklist}
|
||||
|
||||
def _update_models_from_file(self):
|
||||
"""
|
||||
从 models.yaml 和 models.private.yml 中更新 MODEL_NAMES
|
||||
"""
|
||||
|
||||
with open(Path("src/static/models.yaml"), encoding="utf-8") as f:
|
||||
_models = yaml.safe_load(f)
|
||||
|
||||
# 尝试打开一个 models.private.yml 文件，用来覆盖 models.yaml 中的配置
|
||||
try:
|
||||
with open(Path("src/static/models.private.yml"), encoding="utf-8") as f:
|
||||
_models_private = yaml.safe_load(f)
|
||||
except FileNotFoundError:
|
||||
_models_private = {}
|
||||
|
||||
# 修改为按照子元素合并
|
||||
# _models = {**_models, **_models_private}
|
||||
|
||||
self.model_names = {**_models["MODEL_NAMES"], **_models_private.get("MODEL_NAMES", {})}
|
||||
self.embed_model_names = {**_models["EMBED_MODEL_INFO"], **_models_private.get("EMBED_MODEL_INFO", {})}
|
||||
self.reranker_names = {**_models["RERANKER_LIST"], **_models_private.get("RERANKER_LIST", {})}
|
||||
|
||||
def _save_models_to_file(self):
|
||||
_models = {
|
||||
"MODEL_NAMES": self.model_names,
|
||||
"EMBED_MODEL_INFO": self.embed_model_names,
|
||||
"RERANKER_LIST": self.reranker_names,
|
||||
}
|
||||
with open(Path("src/static/models.private.yml"), "w", encoding="utf-8") as f:
|
||||
yaml.dump(_models, f, indent=2, allow_unicode=True)
|
||||
|
||||
def handle_self(self):
|
||||
"""
|
||||
处理配置
|
||||
"""
|
||||
self.model_dir = os.environ.get("MODEL_DIR", "")
|
||||
|
||||
if self.model_dir:
|
||||
if os.path.exists(self.model_dir):
|
||||
logger.debug(
|
||||
f"The model directory （{self.model_dir}） "
|
||||
f"contains the following folders: {os.listdir(self.model_dir)}"
|
||||
)
|
||||
else:
|
||||
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"
|
||||
)
|
||||
|
||||
# 检查模型提供商的环境变量
|
||||
conds = {}
|
||||
self.model_provider_status = {}
|
||||
for provider in self.model_names:
|
||||
conds[provider] = self.model_names[provider]["env"]
|
||||
conds_bool = [bool(os.getenv(_k)) for _k in conds[provider]]
|
||||
self.model_provider_status[provider] = all(conds_bool)
|
||||
|
||||
if os.getenv("TAVILY_API_KEY"):
|
||||
self.enable_web_search = True
|
||||
|
||||
self.valuable_model_provider = [k for k, v in self.model_provider_status.items() if v]
|
||||
assert len(self.valuable_model_provider) > 0, (
|
||||
f"No model provider available, please check your `.env` file. API_KEY_LIST: {conds}"
|
||||
)
|
||||
|
||||
def load(self):
|
||||
"""根据传入的文件覆盖掉默认配置"""
|
||||
logger.info(f"Loading config from {self.filename}")
|
||||
if self.filename is not None and os.path.exists(self.filename):
|
||||
if self.filename.endswith(".json"):
|
||||
with open(self.filename) as f:
|
||||
content = f.read()
|
||||
if content:
|
||||
local_config = json.loads(content)
|
||||
self.update(local_config)
|
||||
else:
|
||||
print(f"{self.filename} is empty.")
|
||||
|
||||
elif self.filename.endswith(".yaml"):
|
||||
with open(self.filename) as f:
|
||||
content = f.read()
|
||||
if content:
|
||||
local_config = yaml.safe_load(content)
|
||||
self.update(local_config)
|
||||
else:
|
||||
print(f"{self.filename} is empty.")
|
||||
else:
|
||||
logger.warning(f"Unknown config file type {self.filename}")
|
||||
|
||||
def save(self):
|
||||
logger.info(f"Saving config to {self.filename}")
|
||||
if self.filename is None:
|
||||
logger.warning("Config file is not specified, save to default config/base.yaml")
|
||||
self.filename = os.path.join(self.save_dir, "config", "base.yaml")
|
||||
os.makedirs(os.path.dirname(self.filename), exist_ok=True)
|
||||
|
||||
if self.filename.endswith(".json"):
|
||||
with open(self.filename, "w+") as f:
|
||||
json.dump(self.__dict__(), f, indent=4, ensure_ascii=False)
|
||||
elif self.filename.endswith(".yaml"):
|
||||
with open(self.filename, "w+") as f:
|
||||
yaml.dump(self.__dict__(), f, indent=2, allow_unicode=True)
|
||||
else:
|
||||
logger.warning(f"Unknown config file type {self.filename}, save as json")
|
||||
with open(self.filename, "w+") as f:
|
||||
json.dump(self, f, indent=4)
|
||||
|
||||
logger.info(f"Config file {self.filename} saved")
|
||||
|
||||
def dump_config(self):
|
||||
return json.loads(str(self))
|
||||
|
||||
def compare_custom_models(self, value):
|
||||
"""
|
||||
比较 custom_models 中的 api_key，如果输入的 api_key 与当前的 api_key 相同，则不修改
|
||||
如果输入的 api_key 为 DEFAULT_MOCK_API，则使用当前的 api_key
|
||||
"""
|
||||
current_models_dict = {model["custom_id"]: model.get("api_key") for model in self.get("custom_models", [])}
|
||||
|
||||
for i, model in enumerate(value):
|
||||
input_custom_id = model.get("custom_id")
|
||||
input_api_key = model.get("api_key")
|
||||
|
||||
if input_custom_id in current_models_dict:
|
||||
current_api_key = current_models_dict[input_custom_id]
|
||||
if input_api_key == DEFAULT_MOCK_API or input_api_key == current_api_key:
|
||||
value[i]["api_key"] = current_api_key
|
||||
|
||||
return value
|
||||
__all__ = ["config"]
|
||||
231
src/config/app.py
Normal file
231
src/config/app.py
Normal file
@ -0,0 +1,231 @@
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
from src.utils.logging_config import logger
|
||||
|
||||
DEFAULT_MOCK_API = "this_is_mock_api_key_in_frontend"
|
||||
|
||||
|
||||
class SimpleConfig(dict):
|
||||
def __key(self, key):
|
||||
return "" if key is None else key # 目前忘记了这里为什么要 lower 了，只能说配置项最好不要有大写的
|
||||
|
||||
def __str__(self):
|
||||
return json.dumps(self)
|
||||
|
||||
def __setattr__(self, key, value):
|
||||
self[self.__key(key)] = value
|
||||
|
||||
def __getattr__(self, key):
|
||||
return self.get(self.__key(key))
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self.get(self.__key(key))
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
return super().__setitem__(self.__key(key), value)
|
||||
|
||||
def __dict__(self):
|
||||
return {k: v for k, v in self.items()}
|
||||
|
||||
def update(self, other):
|
||||
for key, value in other.items():
|
||||
self[key] = value
|
||||
|
||||
|
||||
class Config(SimpleConfig):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._config_items = {}
|
||||
self.save_dir = os.getenv("SAVE_DIR", "saves")
|
||||
self.filename = str(Path(f"{self.save_dir}/config/base.yaml"))
|
||||
os.makedirs(os.path.dirname(self.filename), exist_ok=True)
|
||||
|
||||
self._update_models_from_file()
|
||||
|
||||
### >>> 默认配置
|
||||
# 功能选项
|
||||
self.add_item("enable_reranker", default=False, des="是否开启重排序")
|
||||
self.add_item(
|
||||
"enable_web_search",
|
||||
default=False,
|
||||
des=(
|
||||
"是否开启网页搜索（注：现阶段会根据 TAVILY_API_KEY 自动开启，无法手动配置，将会在下个版本移除此配置项）"
|
||||
),
|
||||
)
|
||||
# 默认智能体配置
|
||||
self.add_item("default_agent_id", default="", des="默认智能体ID")
|
||||
# 模型配置
|
||||
## 注意这里是模型名，而不是具体的模型路径，默认使用 HuggingFace 的路径
|
||||
## 如果需要自定义本地模型路径，则在 src/.env 中配置 MODEL_DIR
|
||||
self.add_item("model_provider", default="siliconflow", des="模型提供商", choices=list(self.model_names.keys()))
|
||||
self.add_item("model_name", default="zai-org/GLM-4.5", des="模型名称")
|
||||
|
||||
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.load()
|
||||
self.handle_self()
|
||||
|
||||
def add_item(self, key, default, des=None, choices=None):
|
||||
self.__setattr__(key, default)
|
||||
self._config_items[key] = {"default": default, "des": des, "choices": choices}
|
||||
|
||||
def __dict__(self):
|
||||
blocklist = [
|
||||
"_config_items",
|
||||
"model_names",
|
||||
"model_provider_status",
|
||||
"embed_model_names",
|
||||
"reranker_names",
|
||||
]
|
||||
return {k: v for k, v in self.items() if k not in blocklist}
|
||||
|
||||
def _update_models_from_file(self):
|
||||
"""
|
||||
从 models.yaml 和 models.private.yml 中更新 MODEL_NAMES
|
||||
"""
|
||||
|
||||
with open(Path("src/static/models.yaml"), encoding="utf-8") as f:
|
||||
_models = yaml.safe_load(f)
|
||||
|
||||
# 尝试打开一个 models.private.yml 文件，用来覆盖 models.yaml 中的配置
|
||||
try:
|
||||
with open(Path("src/static/models.private.yml"), encoding="utf-8") as f:
|
||||
_models_private = yaml.safe_load(f)
|
||||
except FileNotFoundError:
|
||||
_models_private = {}
|
||||
|
||||
# 修改为按照子元素合并
|
||||
# _models = {**_models, **_models_private}
|
||||
|
||||
self.model_names = {**_models["MODEL_NAMES"], **_models_private.get("MODEL_NAMES", {})}
|
||||
self.embed_model_names = {**_models["EMBED_MODEL_INFO"], **_models_private.get("EMBED_MODEL_INFO", {})}
|
||||
self.reranker_names = {**_models["RERANKER_LIST"], **_models_private.get("RERANKER_LIST", {})}
|
||||
|
||||
def _save_models_to_file(self):
|
||||
_models = {
|
||||
"MODEL_NAMES": self.model_names,
|
||||
"EMBED_MODEL_INFO": self.embed_model_names,
|
||||
"RERANKER_LIST": self.reranker_names,
|
||||
}
|
||||
with open(Path("src/static/models.private.yml"), "w", encoding="utf-8") as f:
|
||||
yaml.dump(_models, f, indent=2, allow_unicode=True)
|
||||
|
||||
def handle_self(self):
|
||||
"""
|
||||
处理配置
|
||||
"""
|
||||
self.model_dir = os.environ.get("MODEL_DIR", "")
|
||||
|
||||
if self.model_dir:
|
||||
if os.path.exists(self.model_dir):
|
||||
logger.debug(
|
||||
f"The model directory （{self.model_dir}） "
|
||||
f"contains the following folders: {os.listdir(self.model_dir)}"
|
||||
)
|
||||
else:
|
||||
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"
|
||||
)
|
||||
|
||||
# 检查模型提供商的环境变量
|
||||
conds = {}
|
||||
self.model_provider_status = {}
|
||||
for provider in self.model_names:
|
||||
conds[provider] = self.model_names[provider]["env"]
|
||||
conds_bool = [bool(os.getenv(_k)) for _k in conds[provider]]
|
||||
self.model_provider_status[provider] = all(conds_bool)
|
||||
|
||||
if os.getenv("TAVILY_API_KEY"):
|
||||
self.enable_web_search = True
|
||||
|
||||
self.valuable_model_provider = [k for k, v in self.model_provider_status.items() if v]
|
||||
assert len(self.valuable_model_provider) > 0, (
|
||||
f"No model provider available, please check your `.env` file. API_KEY_LIST: {conds}"
|
||||
)
|
||||
|
||||
def load(self):
|
||||
"""根据传入的文件覆盖掉默认配置"""
|
||||
logger.info(f"Loading config from {self.filename}")
|
||||
if self.filename is not None and os.path.exists(self.filename):
|
||||
if self.filename.endswith(".json"):
|
||||
with open(self.filename) as f:
|
||||
content = f.read()
|
||||
if content:
|
||||
local_config = json.loads(content)
|
||||
self.update(local_config)
|
||||
else:
|
||||
print(f"{self.filename} is empty.")
|
||||
|
||||
elif self.filename.endswith(".yaml"):
|
||||
with open(self.filename) as f:
|
||||
content = f.read()
|
||||
if content:
|
||||
local_config = yaml.safe_load(content)
|
||||
self.update(local_config)
|
||||
else:
|
||||
print(f"{self.filename} is empty.")
|
||||
else:
|
||||
logger.warning(f"Unknown config file type {self.filename}")
|
||||
|
||||
def save(self):
|
||||
logger.info(f"Saving config to {self.filename}")
|
||||
if self.filename is None:
|
||||
logger.warning("Config file is not specified, save to default config/base.yaml")
|
||||
self.filename = os.path.join(self.save_dir, "config", "base.yaml")
|
||||
os.makedirs(os.path.dirname(self.filename), exist_ok=True)
|
||||
|
||||
if self.filename.endswith(".json"):
|
||||
with open(self.filename, "w+") as f:
|
||||
json.dump(self.__dict__(), f, indent=4, ensure_ascii=False)
|
||||
elif self.filename.endswith(".yaml"):
|
||||
with open(self.filename, "w+") as f:
|
||||
yaml.dump(self.__dict__(), f, indent=2, allow_unicode=True)
|
||||
else:
|
||||
logger.warning(f"Unknown config file type {self.filename}, save as json")
|
||||
with open(self.filename, "w+") as f:
|
||||
json.dump(self, f, indent=4)
|
||||
|
||||
logger.info(f"Config file {self.filename} saved")
|
||||
|
||||
def dump_config(self):
|
||||
return json.loads(str(self))
|
||||
|
||||
def compare_custom_models(self, value):
|
||||
"""
|
||||
比较 custom_models 中的 api_key，如果输入的 api_key 与当前的 api_key 相同，则不修改
|
||||
如果输入的 api_key 为 DEFAULT_MOCK_API，则使用当前的 api_key
|
||||
"""
|
||||
current_models_dict = {model["custom_id"]: model.get("api_key") for model in self.get("custom_models", [])}
|
||||
|
||||
for i, model in enumerate(value):
|
||||
input_custom_id = model.get("custom_id")
|
||||
input_api_key = model.get("api_key")
|
||||
|
||||
if input_custom_id in current_models_dict:
|
||||
current_api_key = current_models_dict[input_custom_id]
|
||||
if input_api_key == DEFAULT_MOCK_API or input_api_key == current_api_key:
|
||||
value[i]["api_key"] = current_api_key
|
||||
|
||||
return value
|
||||
|
||||
|
||||
config = Config()
|
||||
@ -1,3 +1,24 @@
|
||||
from .graphbase import GraphDatabase
|
||||
import os
|
||||
|
||||
from ..config import config
|
||||
from .graphbase import GraphDatabase
|
||||
from .chroma_kb import ChromaKB
|
||||
from .kb_factory import KnowledgeBaseFactory
|
||||
from .kb_manager import KnowledgeBaseManager
|
||||
from .lightrag_kb import LightRagKB
|
||||
from .milvus_kb import MilvusKB
|
||||
|
||||
# 注册知识库类型
|
||||
KnowledgeBaseFactory.register("chroma", ChromaKB, {"description": "基于 ChromaDB 的轻量级向量知识库，适合开发和小规模"})
|
||||
KnowledgeBaseFactory.register("milvus", MilvusKB, {"description": "基于 Milvus 的生产级向量知识库，适合高性能部署"})
|
||||
KnowledgeBaseFactory.register("lightrag", LightRagKB, {"description": "基于图检索的知识库，支持实体关系构建和复杂查询"})
|
||||
|
||||
# 创建知识库管理器
|
||||
work_dir = os.path.join(config.save_dir, "knowledge_base_data")
|
||||
knowledge_base = KnowledgeBaseManager(work_dir)
|
||||
|
||||
# 创建图数据库实例
|
||||
graph_base = GraphDatabase()
|
||||
|
||||
__all__ = ["GraphDatabase", "knowledge_base", "graph_base"]
|
||||
|
||||
__all__ = ["GraphDatabase"]
|
||||
|
||||
@ -119,7 +119,7 @@ class OllamaEmbedding(BaseEmbeddingModel):
|
||||
raise ValueError(f"Ollama Embedding failed: Invalid response format {result}")
|
||||
return result["embeddings"]
|
||||
except (requests.RequestException, json.JSONDecodeError) as e:
|
||||
logger.error(f"Ollama Embedding request failed: {e}")
|
||||
logger.error(f"Ollama Embedding request failed: {e}, {payload}")
|
||||
raise ValueError(f"Ollama Embedding request failed: {e}")
|
||||
|
||||
async def apredict(self, message: list[str] | str) -> list[list[float]]:
|
||||
@ -136,7 +136,7 @@ class OllamaEmbedding(BaseEmbeddingModel):
|
||||
raise ValueError(f"Ollama Embedding failed: Invalid response format {result}")
|
||||
return result["embeddings"]
|
||||
except (httpx.RequestError, json.JSONDecodeError) as e:
|
||||
logger.error(f"Ollama Embedding async request failed: {e}")
|
||||
logger.error(f"Ollama Embedding async request failed: {e}, {payload}")
|
||||
raise ValueError(f"Ollama Embedding async request failed: {e}")
|
||||
|
||||
|
||||
@ -158,7 +158,7 @@ class OtherEmbedding(BaseEmbeddingModel):
|
||||
raise ValueError(f"Other Embedding failed: Invalid response format {result}")
|
||||
return [item["embedding"] for item in result["data"]]
|
||||
except (requests.RequestException, json.JSONDecodeError) as e:
|
||||
logger.error(f"Other Embedding request failed: {e}")
|
||||
logger.error(f"Other Embedding request failed: {e}, {payload}")
|
||||
raise ValueError(f"Other Embedding request failed: {e}")
|
||||
|
||||
async def apredict(self, message: list[str] | str) -> list[list[float]]:
|
||||
@ -172,5 +172,5 @@ class OtherEmbedding(BaseEmbeddingModel):
|
||||
raise ValueError(f"Other Embedding failed: Invalid response format {result}")
|
||||
return [item["embedding"] for item in result["data"]]
|
||||
except (httpx.RequestError, json.JSONDecodeError) as e:
|
||||
logger.error(f"Other Embedding async request failed: {e}")
|
||||
logger.error(f"Other Embedding async request failed: {e}, {payload}")
|
||||
raise ValueError(f"Other Embedding async request failed: {e}")
|
||||
|
||||
Loading…
Reference in New Issue
Block a user