From cb9edfdb5eab0779409647a946fed10d9a5d5b21 Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Tue, 16 Dec 2025 00:10:05 +0800 Subject: [PATCH] =?UTF-8?q?feat(graph):=20=E6=94=AF=E6=8C=81=E5=B8=A6?= =?UTF-8?q?=E5=B1=9E=E6=80=A7=E7=9A=84=E7=9F=A5=E8=AF=86=E5=9B=BE=E8=B0=B1?= =?UTF-8?q?=E8=8A=82=E7=82=B9=E5=92=8C=E5=85=B3=E7=B3=BB=E5=AF=BC=E5=85=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 扩展三元组导入功能，支持节点和关系的属性存储 - 新增测试数据文件和单元测试验证功能 - 更新文档说明支持新旧两种数据格式 - 优化查询结果处理，保留并返回节点和关系的属性 --- docs/latest/intro/knowledge-base.md | 26 ++++- src/knowledge/graph.py | 160 ++++++++++++++++++++-------- test/data/complex_graph_test.jsonl | 6 ++ test/test_graph_unit.py | 96 +++++++++++++++++ 4 files changed, 242 insertions(+), 46 deletions(-) create mode 100644 test/data/complex_graph_test.jsonl create mode 100644 test/test_graph_unit.py diff --git a/docs/latest/intro/knowledge-base.md b/docs/latest/intro/knowledge-base.md index d0621ee5..995501fa 100644 --- a/docs/latest/intro/knowledge-base.md +++ b/docs/latest/intro/knowledge-base.md @@ -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 实例 diff --git a/src/knowledge/graph.py b/src/knowledge/graph.py index 61e2c85a..9d488f5f 100644 --- a/src/knowledge/graph.py +++ b/src/knowledge/graph.py @@ -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 diff --git a/test/data/complex_graph_test.jsonl b/test/data/complex_graph_test.jsonl new file mode 100644 index 00000000..d36b0de9 --- /dev/null +++ b/test/data/complex_graph_test.jsonl @@ -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年"}} \ No newline at end of file diff --git a/test/test_graph_unit.py b/test/test_graph_unit.py new file mode 100644 index 00000000..0afceb5e --- /dev/null +++ b/test/test_graph_unit.py @@ -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!")