feat(graph): 支持带属性的知识图谱节点和关系导入

- 扩展三元组导入功能,支持节点和关系的属性存储
- 新增测试数据文件和单元测试验证功能
- 更新文档说明支持新旧两种数据格式
- 优化查询结果处理,保留并返回节点和关系的属性
This commit is contained in:
Wenjie Zhang 2025-12-16 00:10:05 +08:00
parent 3d63418adf
commit cb9edfdb5e
4 changed files with 242 additions and 46 deletions

View File

@ -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 实例

View File

@ -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

View 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
View 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!")