feat: 更新环境模板，添加知识图谱模型名称配置；重构工具工厂，移除冗余代码并优化工具注册逻辑；增强图数据库类，添加数据库连接检查和嵌入模型获取功能；调整设置视图，移除不必要的开关和选择项
This commit is contained in:
Wenjie Zhang
parent
6810f1f84e
commit
9a162d5197
@ -13,3 +13,6 @@ TOGETHER_API_KEY=tgp_v1_fPjW******irD6zesAn4
|
||||
|
||||
# 功能服务
|
||||
TAVILY_API_KEY=tvly-3gR4ind9******JMOxw3E2LG # <<< 配置网络搜索
|
||||
|
||||
# 知识图谱的模型，格式参考 models.yaml
|
||||
GRAPH_EMBED_MODEL_NAME=""
|
||||
@ -1,70 +1,14 @@
|
||||
import json
|
||||
import re
|
||||
import asyncio
|
||||
from collections.abc import Callable
|
||||
from typing import Annotated, Any
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from langchain_tavily import TavilySearch
|
||||
from langchain_core.tools import BaseTool, StructuredTool, tool
|
||||
from pydantic import BaseModel, Field
|
||||
from langchain_core.tools import StructuredTool, tool
|
||||
from langchain_tavily import TavilySearch
|
||||
|
||||
from src import config, graph_base, knowledge_base
|
||||
|
||||
|
||||
# refs https://github.com/chatchat-space/LangGraph-Chatchat chatchat-server/chatchat/server/agent/tools_factory/tools_registry.py
|
||||
def regist_tool(
|
||||
*args: Any,
|
||||
title: str = "",
|
||||
description: str = "",
|
||||
return_direct: bool = False,
|
||||
args_schema: type[BaseModel] | None = None,
|
||||
infer_schema: bool = True,
|
||||
) -> Callable | BaseTool:
|
||||
"""
|
||||
wrapper of langchain tool decorator
|
||||
add tool to registry automatically
|
||||
"""
|
||||
|
||||
def _parse_tool(t: BaseTool):
|
||||
nonlocal description, title
|
||||
|
||||
_TOOLS_REGISTRY[t.name] = t
|
||||
|
||||
# change default description
|
||||
if not description:
|
||||
if t.func is not None:
|
||||
description = t.func.__doc__
|
||||
elif t.coroutine is not None:
|
||||
description = t.coroutine.__doc__
|
||||
t.description = " ".join(re.split(r"\n+\s*", description))
|
||||
# set a default title for human
|
||||
if not title:
|
||||
title = "".join([x.capitalize() for x in t.name.split("_")])
|
||||
setattr(t, "_title", title)
|
||||
|
||||
def wrapper(def_func: Callable) -> BaseTool:
|
||||
partial_ = tool(
|
||||
*args,
|
||||
return_direct=return_direct,
|
||||
args_schema=args_schema,
|
||||
infer_schema=infer_schema,
|
||||
)
|
||||
t = partial_(def_func)
|
||||
_parse_tool(t)
|
||||
return t
|
||||
|
||||
if len(args) == 0:
|
||||
return wrapper
|
||||
else:
|
||||
t = tool(
|
||||
*args,
|
||||
return_direct=return_direct,
|
||||
args_schema=args_schema,
|
||||
infer_schema=infer_schema,
|
||||
)
|
||||
_parse_tool(t)
|
||||
return t
|
||||
from src.utils import logger
|
||||
|
||||
|
||||
class KnowledgeRetrieverModel(BaseModel):
|
||||
@ -75,8 +19,6 @@ class KnowledgeRetrieverModel(BaseModel):
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
|
||||
def get_all_tools():
|
||||
"""获取所有工具"""
|
||||
tools = _TOOLS_REGISTRY.copy()
|
||||
@ -124,7 +66,7 @@ class BaseToolOutput:
|
||||
def __init__(
|
||||
self,
|
||||
data: Any,
|
||||
format: str | Callable = None,
|
||||
format: str | Callable | None = None,
|
||||
data_alias: str = "",
|
||||
**extras: Any,
|
||||
) -> None:
|
||||
|
||||
@ -4,8 +4,10 @@ import warnings
|
||||
import traceback
|
||||
|
||||
from neo4j import GraphDatabase as GD
|
||||
from neo4j import Query
|
||||
|
||||
from src import config
|
||||
from src.models.embedding import get_embedding_model
|
||||
from src.utils import logger
|
||||
|
||||
warnings.filterwarnings("ignore", category=UserWarning)
|
||||
@ -19,7 +21,8 @@ class GraphDatabase:
|
||||
self.files = []
|
||||
self.status = "closed"
|
||||
self.kgdb_name = "neo4j"
|
||||
self.embed_model_name = None
|
||||
self.embed_model_name = os.getenv("GRAPH_EMBED_MODEL_NAME") or "siliconflow/BAAI/bge-m3"
|
||||
self.embed_model = get_embedding_model(self.embed_model_name)
|
||||
self.work_dir = os.path.join(config.save_dir, "knowledge_graph", self.kgdb_name)
|
||||
os.makedirs(self.work_dir, exist_ok=True)
|
||||
|
||||
@ -45,6 +48,7 @@ class GraphDatabase:
|
||||
|
||||
def close(self):
|
||||
"""关闭数据库连接"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.driver.close()
|
||||
|
||||
def is_running(self):
|
||||
@ -53,6 +57,7 @@ class GraphDatabase:
|
||||
|
||||
def get_sample_nodes(self, kgdb_name='neo4j', num=50):
|
||||
"""获取指定数据库的 num 个节点信息"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.use_database(kgdb_name)
|
||||
def query(tx, num):
|
||||
result = tx.run("MATCH (n)-[r]->(m) RETURN n, r, m LIMIT $num", num=int(num))
|
||||
@ -63,6 +68,7 @@ class GraphDatabase:
|
||||
|
||||
def create_graph_database(self, kgdb_name):
|
||||
"""创建新的数据库，如果已存在则返回已有数据库的名称"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
with self.driver.session() as session:
|
||||
existing_databases = session.run("SHOW DATABASES")
|
||||
existing_db_names = [db['name'] for db in existing_databases]
|
||||
@ -71,7 +77,7 @@ class GraphDatabase:
|
||||
print(f"已存在数据库: {existing_db_names[0]}")
|
||||
return existing_db_names[0] # 返回所有已有数据库名称
|
||||
|
||||
session.run(f"CREATE DATABASE {kgdb_name}")
|
||||
session.run(f"CREATE DATABASE {kgdb_name}") # type: ignore
|
||||
print(f"数据库 '{kgdb_name}' 创建成功.")
|
||||
return kgdb_name # 返回创建的数据库名称
|
||||
|
||||
@ -83,6 +89,7 @@ class GraphDatabase:
|
||||
|
||||
def txt_add_entity(self, triples, kgdb_name='neo4j'):
|
||||
"""添加实体三元组"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.use_database(kgdb_name)
|
||||
def create(tx, triples):
|
||||
for triple in triples:
|
||||
@ -101,6 +108,7 @@ class GraphDatabase:
|
||||
|
||||
async def txt_add_vector_entity(self, triples, kgdb_name='neo4j'):
|
||||
"""添加实体三元组"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.use_database(kgdb_name)
|
||||
def _index_exists(tx, index_name):
|
||||
"""检查索引是否存在"""
|
||||
@ -211,6 +219,7 @@ class GraphDatabase:
|
||||
self.save_graph_info()
|
||||
|
||||
async def jsonl_file_add_entity(self, file_path, kgdb_name='neo4j'):
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.status = "processing"
|
||||
kgdb_name = kgdb_name or 'neo4j'
|
||||
self.use_database(kgdb_name) # 切换到指定数据库
|
||||
@ -233,6 +242,7 @@ class GraphDatabase:
|
||||
|
||||
def delete_entity(self, entity_name=None, kgdb_name="neo4j"):
|
||||
"""删除数据库中的指定实体三元组, 参数entity_name为空则删除全部实体"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.use_database(kgdb_name)
|
||||
with self.driver.session() as session:
|
||||
if entity_name:
|
||||
@ -256,6 +266,7 @@ class GraphDatabase:
|
||||
|
||||
def query_node(self, entity_name, threshold=0.9, kgdb_name='neo4j', hops=2, max_entities=5, **kwargs):
|
||||
"""知识图谱查询节点的入口:"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
# TODO 添加判断节点数量为 0 停止检索
|
||||
# 判断是否启动
|
||||
if not self.is_running():
|
||||
@ -306,6 +317,7 @@ class GraphDatabase:
|
||||
|
||||
def query_specific_entity(self, entity_name, kgdb_name='neo4j', hops=2, limit=100):
|
||||
"""查询指定实体三元组信息（无向关系）"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
if not entity_name:
|
||||
logger.warning("实体名称为空")
|
||||
return []
|
||||
@ -343,6 +355,7 @@ class GraphDatabase:
|
||||
|
||||
def query_all_nodes_and_relationships(self, kgdb_name='neo4j', hops = 2):
|
||||
"""查询图数据库中所有三元组信息 NEVER USE"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.use_database(kgdb_name)
|
||||
def query(tx, hops):
|
||||
result = tx.run(f"""
|
||||
@ -358,6 +371,7 @@ class GraphDatabase:
|
||||
|
||||
def query_by_relationship_type(self, relationship_type, kgdb_name='neo4j', hops = 2):
|
||||
"""查询指定关系三元组信息 NEVER USE"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.use_database(kgdb_name)
|
||||
def query(tx, relationship_type, hops):
|
||||
result = tx.run(f"""
|
||||
@ -373,6 +387,7 @@ class GraphDatabase:
|
||||
|
||||
def query_entity_like(self, keyword, kgdb_name='neo4j', hops = 2):
|
||||
"""模糊查询 NEVER USE"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.use_database(kgdb_name)
|
||||
def query(tx, keyword, hops):
|
||||
result = tx.run(f"""
|
||||
@ -390,6 +405,7 @@ class GraphDatabase:
|
||||
|
||||
def query_node_info(self, node_name, kgdb_name='neo4j', hops = 2):
|
||||
"""查询指定节点的详细信息返回信息 NEVER USE"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.use_database(kgdb_name) # 切换到指定数据库
|
||||
def query(tx, node_name, hops):
|
||||
result = tx.run(f"""
|
||||
@ -405,23 +421,19 @@ class GraphDatabase:
|
||||
return session.execute_read(query, node_name, hops)
|
||||
|
||||
async def aget_embedding(self, text):
|
||||
from src import knowledge_base
|
||||
|
||||
if isinstance(text, list):
|
||||
outputs = await knowledge_base.embed_model.abatch_encode(text, batch_size=40)
|
||||
outputs = await self.embed_model.abatch_encode(text, batch_size=40)
|
||||
return outputs
|
||||
else:
|
||||
outputs = await knowledge_base.embed_model.aencode(text)
|
||||
outputs = await self.embed_model.aencode(text)
|
||||
return outputs
|
||||
|
||||
def get_embedding(self, text):
|
||||
from src import knowledge_base
|
||||
|
||||
if isinstance(text, list):
|
||||
outputs = knowledge_base.embed_model.batch_encode(text, batch_size=40)
|
||||
outputs = self.embed_model.batch_encode(text, batch_size=40)
|
||||
return outputs
|
||||
else:
|
||||
outputs = knowledge_base.embed_model.encode([text])[0]
|
||||
outputs = self.embed_model.encode([text])[0]
|
||||
return outputs
|
||||
|
||||
def set_embedding(self, tx, entity_name, embedding):
|
||||
@ -431,6 +443,7 @@ class GraphDatabase:
|
||||
""", name=entity_name, embedding=embedding)
|
||||
|
||||
def get_graph_info(self, graph_name="neo4j"):
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.use_database(graph_name)
|
||||
def query(tx):
|
||||
entity_count = tx.run("MATCH (n) RETURN count(n) AS count").single()["count"]
|
||||
@ -493,6 +506,7 @@ class GraphDatabase:
|
||||
Returns:
|
||||
list: 没有嵌入向量的节点列表
|
||||
"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.use_database(kgdb_name)
|
||||
|
||||
def query(tx):
|
||||
@ -543,6 +557,7 @@ class GraphDatabase:
|
||||
Returns:
|
||||
int: 成功添加嵌入向量的节点数量
|
||||
"""
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.use_database(kgdb_name)
|
||||
|
||||
# 如果node_names为None，则获取所有没有嵌入向量的节点
|
||||
@ -575,6 +590,7 @@ class GraphDatabase:
|
||||
source_id = source.element_id
|
||||
target_id = target.element_id
|
||||
|
||||
assert node_dict is not None, "node_dict is required"
|
||||
source_name = node_dict[source_id]["name"] if source_name is None else source_name
|
||||
target_name = node_dict[target_id]["name"] if target_name is None else target_name
|
||||
|
||||
|
||||
@ -43,7 +43,6 @@
|
||||
<a-select style="width: 300px"
|
||||
:value="configStore.config?.reranker"
|
||||
@change="handleChange('reranker', $event)"
|
||||
:disabled="!configStore.config.enable_reranker"
|
||||
>
|
||||
<a-select-option
|
||||
v-for="(name, idx) in items?.reranker.choices" :key="idx"
|
||||
@ -51,25 +50,6 @@
|
||||
</a-select-option>
|
||||
</a-select>
|
||||
</div>
|
||||
<div class="card">
|
||||
<span class="label">{{ items?.enable_reranker.des }}</span>
|
||||
<a-switch
|
||||
:checked="configStore.config.enable_reranker"
|
||||
@change="handleChange('enable_reranker', !configStore.config.enable_reranker)"
|
||||
/>
|
||||
</div>
|
||||
<div class="card card-select">
|
||||
<span class="label">{{ items?.use_rewrite_query.des }}</span>
|
||||
<a-select style="width: 200px"
|
||||
:value="configStore.config?.use_rewrite_query"
|
||||
@change="handleChange('use_rewrite_query', $event)"
|
||||
>
|
||||
<a-select-option
|
||||
v-for="(name, idx) in items?.use_rewrite_query.choices" :key="idx"
|
||||
:value="name">{{ name }}
|
||||
</a-select-option>
|
||||
</a-select>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="setting" v-if="(state.windowWidth <= 520 || state.section === 'model') && userStore.isSuperAdmin">
|
||||
|
||||
Loading…
Reference in New Issue
Block a user