feat(graph): 支持带属性的知识图谱节点和关系导入
- 扩展三元组导入功能,支持节点和关系的属性存储 - 新增测试数据文件和单元测试验证功能 - 更新文档说明支持新旧两种数据格式 - 优化查询结果处理,保留并返回节点和关系的属性
This commit is contained in:
parent
3d63418adf
commit
cb9edfdb5e
@ -97,16 +97,32 @@ LIMIT $num
|
||||
|
||||
### 1. 以三元组形式导入
|
||||
|
||||
系统支持通过网页导入 `jsonl` 格式的知识图谱数据,支持**简单三元组**和**带属性三元组**两种格式。
|
||||
|
||||
系统支持通过网页导入 `jsonl` 格式的知识图谱数据:
|
||||
**简单格式(兼容旧版)**:
|
||||
|
||||
```jsonl
|
||||
{"h": "北京", "t": "中国", "r": "首都"}
|
||||
{"h": "上海", "t": "中国", "r": "直辖市"}
|
||||
{"h": "深圳", "t": "广东", "r": "省会"}
|
||||
```
|
||||
|
||||
**格式说明**,每行一个三元组,系统自动验证数据格式,并自动导入到 Neo4j 数据库,添加 `Upload`、`Entity`、`Relation` 标签,会自动处理重复的三元组。
|
||||
**扩展格式(支持属性)**:
|
||||
|
||||
支持 `h`(头节点)、`t`(尾节点)和 `r`(关系)为对象结构,其中:
|
||||
- 节点对象必须包含 `name` 字段。
|
||||
- 关系对象必须包含 `type` 字段。
|
||||
- 其他字段将作为**属性**存储在 Neo4j 中。
|
||||
|
||||
```jsonl
|
||||
{"h": {"name": "孙悟空", "title": "齐天大圣", "weapon": "如意金箍棒"}, "t": {"name": "唐僧", "species": "人"}, "r": {"type": "徒弟", "order": 1}}
|
||||
{"h": "猪八戒", "t": {"name": "唐僧"}, "r": {"type": "徒弟", "order": 2}}
|
||||
```
|
||||
|
||||
**格式说明**:
|
||||
- 每行一个数据项。
|
||||
- 系统自动验证数据格式,并自动导入到 Neo4j 数据库。
|
||||
- 自动添加 `Upload`、`Entity` 标签(节点)和 `RELATION` 类型(关系)。
|
||||
- 自动处理重复实体和关系,并合并属性。
|
||||
|
||||
Neo4j 访问信息可以参考 `docker-compose.yml` 中配置对应的环境变量来覆盖。
|
||||
|
||||
@ -116,7 +132,9 @@ Neo4j 访问信息可以参考 `docker-compose.yml` 中配置对应的环境变
|
||||
- **连接地址**: bolt://localhost:7687
|
||||
|
||||
::: tip 测试数据
|
||||
可以使用 `test/data/A_Dream_of_Red_Mansions_tiny.jsonl` 文件进行测试导入。
|
||||
可以使用以下文件进行测试导入:
|
||||
- 简单格式:`test/data/A_Dream_of_Red_Mansions_tiny.jsonl`
|
||||
- 扩展属性格式:`test/data/complex_graph_test.jsonl`
|
||||
:::
|
||||
|
||||
### 2. 接入已有 Neo4j 实例
|
||||
|
||||
@ -64,6 +64,22 @@ class GraphDatabase:
|
||||
assert self.driver is not None, "Database is not connected"
|
||||
self.use_database(kgdb_name)
|
||||
|
||||
def _process_record_props(record):
|
||||
"""处理记录中的属性:扁平化 properties 并移除 embedding"""
|
||||
if record is None:
|
||||
return None
|
||||
|
||||
# 复制一份以避免修改原字典
|
||||
data = dict(record)
|
||||
props = data.pop("properties", {}) or {}
|
||||
|
||||
# 移除 embedding
|
||||
if "embedding" in props:
|
||||
del props["embedding"]
|
||||
|
||||
# 合并属性(优先保留原字典中的 id, name, type 等核心字段)
|
||||
return {**props, **data}
|
||||
|
||||
def query(tx, num):
|
||||
"""Note: 使用连通性查询获取集中的节点子图"""
|
||||
# 首先尝试获取一个连通的子图
|
||||
@ -104,17 +120,18 @@ class GraphDatabase:
|
||||
OPTIONAL MATCH (n)-[rel]-(m)
|
||||
WHERE m IN final_nodes AND elementId(n) < elementId(m)
|
||||
RETURN
|
||||
{id: elementId(n), name: n.name} AS h,
|
||||
{id: elementId(n), name: n.name, properties: properties(n)} AS h,
|
||||
CASE WHEN rel IS NOT NULL THEN
|
||||
{
|
||||
id: elementId(rel),
|
||||
type: rel.type,
|
||||
source_id: elementId(startNode(rel)),
|
||||
target_id: elementId(endNode(rel))
|
||||
target_id: elementId(endNode(rel)),
|
||||
properties: properties(rel)
|
||||
}
|
||||
ELSE null END AS r,
|
||||
CASE WHEN m IS NOT NULL THEN
|
||||
{id: elementId(m), name: m.name}
|
||||
{id: elementId(m), name: m.name, properties: properties(m)}
|
||||
ELSE null END AS t
|
||||
"""
|
||||
|
||||
@ -124,7 +141,7 @@ class GraphDatabase:
|
||||
node_ids = set()
|
||||
|
||||
for item in results:
|
||||
h_node = item["h"]
|
||||
h_node = _process_record_props(item["h"])
|
||||
|
||||
# 始终添加头节点
|
||||
if h_node["id"] not in node_ids:
|
||||
@ -133,14 +150,15 @@ class GraphDatabase:
|
||||
|
||||
# 只有当边和尾节点都存在时才处理
|
||||
if item["r"] is not None and item["t"] is not None:
|
||||
t_node = item["t"]
|
||||
t_node = _process_record_props(item["t"])
|
||||
r_edge = _process_record_props(item["r"])
|
||||
|
||||
# 避免重复添加尾节点
|
||||
if t_node["id"] not in node_ids:
|
||||
formatted_results["nodes"].append(t_node)
|
||||
node_ids.add(t_node["id"])
|
||||
|
||||
formatted_results["edges"].append(item["r"])
|
||||
formatted_results["edges"].append(r_edge)
|
||||
|
||||
# 如果连通查询返回的节点数不足,补充更多节点
|
||||
if len(formatted_results["nodes"]) < num:
|
||||
@ -150,14 +168,14 @@ class GraphDatabase:
|
||||
supplement_query = """
|
||||
MATCH (n:Entity)
|
||||
WHERE NOT elementId(n) IN $existing_ids
|
||||
RETURN {id: elementId(n), name: n.name} AS node
|
||||
RETURN {id: elementId(n), name: n.name, properties: properties(n)} AS node
|
||||
LIMIT $count
|
||||
"""
|
||||
|
||||
supplement_results = tx.run(supplement_query, existing_ids=list(node_ids), count=remaining_count)
|
||||
|
||||
for item in supplement_results:
|
||||
node = item["node"]
|
||||
node = _process_record_props(item["node"])
|
||||
formatted_results["nodes"].append(node)
|
||||
node_ids.add(node["id"])
|
||||
|
||||
@ -170,14 +188,15 @@ class GraphDatabase:
|
||||
MATCH (n:Entity)-[r]-(m:Entity)
|
||||
WHERE elementId(n) < elementId(m)
|
||||
RETURN
|
||||
{id: elementId(n), name: n.name} AS h,
|
||||
{id: elementId(n), name: n.name, properties: properties(n)} AS h,
|
||||
{
|
||||
id: elementId(r),
|
||||
type: r.type,
|
||||
source_id: elementId(startNode(r)),
|
||||
target_id: elementId(endNode(r))
|
||||
target_id: elementId(endNode(r)),
|
||||
properties: properties(r)
|
||||
} AS r,
|
||||
{id: elementId(m), name: m.name} AS t
|
||||
{id: elementId(m), name: m.name, properties: properties(m)} AS t
|
||||
LIMIT $num
|
||||
"""
|
||||
results = tx.run(fallback_query, num=int(num))
|
||||
@ -185,8 +204,9 @@ class GraphDatabase:
|
||||
node_ids = set()
|
||||
|
||||
for item in results:
|
||||
h_node = item["h"]
|
||||
t_node = item["t"]
|
||||
h_node = _process_record_props(item["h"])
|
||||
t_node = _process_record_props(item["t"])
|
||||
r_edge = _process_record_props(item["r"])
|
||||
|
||||
# 避免重复添加节点
|
||||
if h_node["id"] not in node_ids:
|
||||
@ -196,7 +216,7 @@ class GraphDatabase:
|
||||
formatted_results["nodes"].append(t_node)
|
||||
node_ids.add(t_node["id"])
|
||||
|
||||
formatted_results["edges"].append(item["r"])
|
||||
formatted_results["edges"].append(r_edge)
|
||||
|
||||
return formatted_results
|
||||
|
||||
@ -240,18 +260,47 @@ class GraphDatabase:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _parse_node(node_data):
|
||||
"""解析节点数据,返回 (name, props)"""
|
||||
if isinstance(node_data, dict):
|
||||
props = node_data.copy()
|
||||
name = props.pop("name", "")
|
||||
return name, props
|
||||
return str(node_data), {}
|
||||
|
||||
def _parse_relation(rel_data):
|
||||
"""解析关系数据,返回 (type, props)"""
|
||||
if isinstance(rel_data, dict):
|
||||
props = rel_data.copy()
|
||||
rel_type = props.pop("type", "")
|
||||
return rel_type, props
|
||||
return str(rel_data), {}
|
||||
|
||||
def _create_graph(tx, data):
|
||||
"""添加一个三元组"""
|
||||
for entry in data:
|
||||
h_name, h_props = _parse_node(entry.get("h"))
|
||||
t_name, t_props = _parse_node(entry.get("t"))
|
||||
r_type, r_props = _parse_relation(entry.get("r"))
|
||||
|
||||
if not h_name or not t_name or not r_type:
|
||||
continue
|
||||
|
||||
tx.run(
|
||||
"""
|
||||
MERGE (h:Entity:Upload {name: $h})
|
||||
MERGE (t:Entity:Upload {name: $t})
|
||||
MERGE (h)-[r:RELATION {type: $r}]->(t)
|
||||
MERGE (h:Entity:Upload {name: $h_name})
|
||||
SET h += $h_props
|
||||
MERGE (t:Entity:Upload {name: $t_name})
|
||||
SET t += $t_props
|
||||
MERGE (h)-[r:RELATION {type: $r_type}]->(t)
|
||||
SET r += $r_props
|
||||
""",
|
||||
h=entry["h"],
|
||||
t=entry["t"],
|
||||
r=entry["r"],
|
||||
h_name=h_name,
|
||||
h_props=h_props,
|
||||
t_name=t_name,
|
||||
t_props=t_props,
|
||||
r_type=r_type,
|
||||
r_props=r_props,
|
||||
)
|
||||
|
||||
def _create_vector_index(tx, dim):
|
||||
@ -273,6 +322,9 @@ class GraphDatabase:
|
||||
# 构建参数字典,将列表转换为"param0"、"param1"等键值对形式
|
||||
params = {f"param{i}": name for i, name in enumerate(entity_names)}
|
||||
|
||||
if not params:
|
||||
return []
|
||||
|
||||
# 构建查询参数列表
|
||||
param_placeholders = ", ".join([f"${key}" for key in params.keys()])
|
||||
|
||||
@ -315,20 +367,24 @@ class GraphDatabase:
|
||||
session.execute_write(_create_vector_index, getattr(cur_embed_info, "dimension", 1024))
|
||||
|
||||
# 收集所有需要处理的实体名称,去重
|
||||
all_entities = []
|
||||
all_entities = set()
|
||||
for entry in triples:
|
||||
if entry["h"] not in all_entities:
|
||||
all_entities.append(entry["h"])
|
||||
if entry["t"] not in all_entities:
|
||||
all_entities.append(entry["t"])
|
||||
h_name, _ = _parse_node(entry.get("h"))
|
||||
t_name, _ = _parse_node(entry.get("t"))
|
||||
if h_name:
|
||||
all_entities.add(h_name)
|
||||
if t_name:
|
||||
all_entities.add(t_name)
|
||||
|
||||
all_entities_list = list(all_entities)
|
||||
|
||||
# 筛选出没有embedding的节点
|
||||
nodes_without_embedding = session.execute_read(_get_nodes_without_embedding, all_entities)
|
||||
nodes_without_embedding = session.execute_read(_get_nodes_without_embedding, all_entities_list)
|
||||
if not nodes_without_embedding:
|
||||
logger.info("所有实体已有embedding,无需重新计算")
|
||||
return
|
||||
|
||||
logger.info(f"需要为{len(nodes_without_embedding)}/{len(all_entities)}个实体计算embedding")
|
||||
logger.info(f"需要为{len(nodes_without_embedding)}/{len(all_entities_list)}个实体计算embedding")
|
||||
|
||||
# 批量处理实体
|
||||
max_batch_size = 1024 # 限制此部分的主要是内存大小 1024 * 1024 * 4 / 1024 / 1024 = 4GB
|
||||
@ -597,30 +653,46 @@ class GraphDatabase:
|
||||
|
||||
self.use_database(kgdb_name)
|
||||
|
||||
def _process_record_props(record):
|
||||
"""处理记录中的属性:扁平化 properties 并移除 embedding"""
|
||||
if record is None:
|
||||
return None
|
||||
|
||||
# 复制一份以避免修改原字典
|
||||
data = dict(record)
|
||||
props = data.pop("properties", {}) or {}
|
||||
|
||||
# 移除 embedding
|
||||
if "embedding" in props:
|
||||
del props["embedding"]
|
||||
|
||||
# 合并属性(优先保留原字典中的 id, name, type 等核心字段)
|
||||
return {**props, **data}
|
||||
|
||||
def query(tx, entity_name, hops, limit):
|
||||
try:
|
||||
query_str = """
|
||||
WITH [
|
||||
// 1跳出边
|
||||
[(n {name: $entity_name})-[r1]->(m1) |
|
||||
{h: {id: elementId(n), name: n.name},
|
||||
r: {id: elementId(r1), type: r1.type, source_id: elementId(n), target_id: elementId(m1)},
|
||||
t: {id: elementId(m1), name: m1.name}}],
|
||||
{h: {id: elementId(n), name: n.name, properties: properties(n)},
|
||||
r: {id: elementId(r1), type: r1.type, source_id: elementId(n), target_id: elementId(m1), properties: properties(r1)},
|
||||
t: {id: elementId(m1), name: m1.name, properties: properties(m1)}}],
|
||||
// 2跳出边
|
||||
[(n {name: $entity_name})-[r1]->(m1)-[r2]->(m2) |
|
||||
{h: {id: elementId(m1), name: m1.name},
|
||||
r: {id: elementId(r2), type: r2.type, source_id: elementId(m1), target_id: elementId(m2)},
|
||||
t: {id: elementId(m2), name: m2.name}}],
|
||||
{h: {id: elementId(m1), name: m1.name, properties: properties(m1)},
|
||||
r: {id: elementId(r2), type: r2.type, source_id: elementId(m1), target_id: elementId(m2), properties: properties(r2)},
|
||||
t: {id: elementId(m2), name: m2.name, properties: properties(m2)}}],
|
||||
// 1跳入边
|
||||
[(m1)-[r1]->(n {name: $entity_name}) |
|
||||
{h: {id: elementId(m1), name: m1.name},
|
||||
r: {id: elementId(r1), type: r1.type, source_id: elementId(m1), target_id: elementId(n)},
|
||||
t: {id: elementId(n), name: n.name}}],
|
||||
{h: {id: elementId(m1), name: m1.name, properties: properties(m1)},
|
||||
r: {id: elementId(r1), type: r1.type, source_id: elementId(m1), target_id: elementId(n), properties: properties(r1)},
|
||||
t: {id: elementId(n), name: n.name, properties: properties(n)}}],
|
||||
// 2跳入边
|
||||
[(m2)-[r2]->(m1)-[r1]->(n {name: $entity_name}) |
|
||||
{h: {id: elementId(m2), name: m2.name},
|
||||
r: {id: elementId(r2), type: r2.type, source_id: elementId(m2), target_id: elementId(m1)},
|
||||
t: {id: elementId(m1), name: m1.name}}]
|
||||
{h: {id: elementId(m2), name: m2.name, properties: properties(m2)},
|
||||
r: {id: elementId(r2), type: r2.type, source_id: elementId(m2), target_id: elementId(m1), properties: properties(r2)},
|
||||
t: {id: elementId(m1), name: m1.name, properties: properties(m1)}}]
|
||||
] AS all_results
|
||||
UNWIND all_results AS result_list
|
||||
UNWIND result_list AS item
|
||||
@ -636,9 +708,13 @@ class GraphDatabase:
|
||||
formatted_results = {"nodes": [], "edges": [], "triples": []}
|
||||
|
||||
for item in results:
|
||||
formatted_results["nodes"].extend([item["h"], item["t"]])
|
||||
formatted_results["edges"].append(item["r"])
|
||||
formatted_results["triples"].append((item["h"]["name"], item["r"]["type"], item["t"]["name"]))
|
||||
h = _process_record_props(item["h"])
|
||||
r = _process_record_props(item["r"])
|
||||
t = _process_record_props(item["t"])
|
||||
|
||||
formatted_results["nodes"].extend([h, t])
|
||||
formatted_results["edges"].append(r)
|
||||
formatted_results["triples"].append((h["name"], r["type"], t["name"]))
|
||||
|
||||
logger.debug(f"Query Results: {results}")
|
||||
return formatted_results
|
||||
|
||||
6
test/data/complex_graph_test.jsonl
Normal file
6
test/data/complex_graph_test.jsonl
Normal file
@ -0,0 +1,6 @@
|
||||
{"h": {"name": "孙悟空", "title": "齐天大圣", "weapon": "如意金箍棒", "species": "猴"}, "t": {"name": "唐僧", "title": "旃檀功德佛", "species": "人"}, "r": {"type": "徒弟", "order": 1}}
|
||||
{"h": {"name": "猪八戒", "title": "天蓬元帅", "weapon": "九齿钉耙", "species": "猪"}, "t": {"name": "唐僧"}, "r": {"type": "徒弟", "order": 2}}
|
||||
{"h": {"name": "沙悟净", "title": "卷帘大将", "weapon": "降妖宝杖"}, "t": {"name": "唐僧"}, "r": {"type": "徒弟", "order": 3}}
|
||||
{"h": {"name": "白龙马", "origin": "西海龙宫"}, "t": {"name": "唐僧"}, "r": {"type": "坐骑", "original_form": "龙"}}
|
||||
{"h": "孙悟空", "t": {"name": "猪八戒"}, "r": {"type": "师兄弟", "relationship": "conflict/cooperate"}}
|
||||
{"h": {"name": "如来佛祖", "location": "西天雷音寺"}, "t": {"name": "孙悟空"}, "r": {"type": "压制", "tool": "五指山", "duration": "500年"}}
|
||||
96
test/test_graph_unit.py
Normal file
96
test/test_graph_unit.py
Normal file
@ -0,0 +1,96 @@
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, AsyncMock, patch
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Add project root to path
|
||||
sys.path.append(os.getcwd())
|
||||
|
||||
from src.knowledge.graph import GraphDatabase
|
||||
from src import config
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_txt_add_vector_entity_parsing():
|
||||
# Mock driver and session
|
||||
mock_driver = MagicMock()
|
||||
mock_session = MagicMock()
|
||||
mock_driver.session.return_value.__enter__.return_value = mock_session
|
||||
|
||||
# Setup mock transaction
|
||||
mock_tx = MagicMock()
|
||||
def side_effect_execute_write(func, *args, **kwargs):
|
||||
return func(mock_tx, *args, **kwargs)
|
||||
|
||||
# Mock execute_read to return empty list (no missing embeddings for this test)
|
||||
# The code calls _get_nodes_without_embedding which returns [record['name']]
|
||||
# If we return [], it means all nodes have embeddings or none found.
|
||||
# Actually, the code checks:
|
||||
# nodes_without_embedding = session.execute_read(_get_nodes_without_embedding, all_entities)
|
||||
# Let's mock it to return empty list so we skip embedding generation loop which simplifies test
|
||||
mock_session.execute_read.return_value = []
|
||||
|
||||
mock_session.execute_write.side_effect = side_effect_execute_write
|
||||
|
||||
# Mock embedding model
|
||||
with patch('src.knowledge.graph.select_embedding_model') as mock_select_model:
|
||||
mock_embed_model = MagicMock()
|
||||
mock_select_model.return_value = mock_embed_model
|
||||
|
||||
# Instantiate GraphDatabase with mocked driver
|
||||
# We also need to patch GD.driver in the init
|
||||
with patch('src.knowledge.graph.GD.driver', return_value=mock_driver):
|
||||
gd = GraphDatabase()
|
||||
# Manually set driver and status just in case init didn't work as expected due to other mocks
|
||||
gd.driver = mock_driver
|
||||
gd.status = "open"
|
||||
gd.embed_model_name = "test_model" # avoid config check issues if possible
|
||||
|
||||
# Mock config to match
|
||||
with patch('src.knowledge.graph.config.embed_model', "test_model"):
|
||||
with patch('src.knowledge.graph.config.embed_model_names', {"test_model": MagicMock(dimension=1024)}):
|
||||
|
||||
# Test data: Mixed format
|
||||
triples = [
|
||||
# Legacy format
|
||||
{"h": "A", "r": "KNOWS", "t": "B"},
|
||||
# Extended format
|
||||
{
|
||||
"h": {"name": "C", "age": 30},
|
||||
"r": {"type": "LIKES", "weight": 0.8},
|
||||
"t": {"name": "D", "role": "User"}
|
||||
}
|
||||
]
|
||||
|
||||
# Run the method
|
||||
await gd.txt_add_vector_entity(triples)
|
||||
|
||||
# Verify calls to mock_tx.run
|
||||
merge_calls = []
|
||||
for call in mock_tx.run.call_args_list:
|
||||
args, kwargs = call
|
||||
query = args[0] if args else kwargs.get('query', '')
|
||||
if "MERGE (h:Entity:Upload" in query:
|
||||
# The args are passed as kwargs to run: h_name=..., etc.
|
||||
merge_calls.append(kwargs)
|
||||
|
||||
assert len(merge_calls) == 2, f"Expected 2 merge calls, got {len(merge_calls)}"
|
||||
|
||||
# Call 1 (Legacy)
|
||||
call1 = merge_calls[0]
|
||||
assert call1['h_name'] == "A"
|
||||
assert call1['h_props'] == {}
|
||||
assert call1['t_name'] == "B"
|
||||
assert call1['t_props'] == {}
|
||||
assert call1['r_type'] == "KNOWS"
|
||||
assert call1['r_props'] == {}
|
||||
|
||||
# Call 2 (Extended)
|
||||
call2 = merge_calls[1]
|
||||
assert call2['h_name'] == "C"
|
||||
assert call2['h_props'] == {'age': 30}
|
||||
assert call2['t_name'] == "D"
|
||||
assert call2['t_props'] == {'role': 'User'}
|
||||
assert call2['r_type'] == "LIKES"
|
||||
assert call2['r_props'] == {'weight': 0.8}
|
||||
|
||||
print("Verification passed!")
|
||||
Loading…
Reference in New Issue
Block a user