gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix linting

Browse Source
This commit is contained in:
yangdx
2025-04-16 14:07:22 +08:00
parent e6b2a035ea
commit 2a950f3ff9
2 changed files with 146 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
lightrag/kg/postgres_impl.py
View File

@@ -1037,7 +1037,10 @@ class PGGraphStorage(BaseGraphStorage):
await self.upsert_node(node1_id, node1_data)
await self.delete_node(node1_id)
query = """CREATE INDEX entity_id_gin_idxSELECT ON %s."base" USING gin (properties);""" % (self.graph_name)
query = (
"""CREATE INDEX entity_id_gin_idxSELECT ON %s."base" USING gin (properties);"""
% (self.graph_name)
)
await self.db.execute(
query,
@@ -1659,11 +1662,15 @@ class PGGraphStorage(BaseGraphStorage):
for result in forward_results:
if result["source"] and result["target"] and result["edge_properties"]:
edges_dict[(result["source"], result["target"])] = result["edge_properties"]
edges_dict[(result["source"], result["target"])] = result[
"edge_properties"
]
for result in backward_results:
if result["source"] and result["target"] and result["edge_properties"]:
edges_dict[(result["source"], result["target"])] = result["edge_properties"]
edges_dict[(result["source"], result["target"])] = result[
"edge_properties"
]
return edges_dict

114
tests/test_graph_storage.py
View File

@@ -219,11 +219,15 @@ async def test_graph_basic(storage):
print(f"反向边描述: {reverse_edge_props.get('description', '无描述')}")
print(f"反向边权重: {reverse_edge_props.get('weight', '无权重')}")
# 验证正向和反向边属性是否相同
assert edge_props == reverse_edge_props, "正向和反向边属性不一致，无向图特性验证失败"
assert (
edge_props == reverse_edge_props
), "正向和反向边属性不一致，无向图特性验证失败"
print("无向图特性验证成功：正向和反向边属性一致")
else:
print(f"读取反向边属性失败: {node2_id} -> {node1_id}")
assert False, f"未能读取反向边属性: {node2_id} -> {node1_id}，无向图特性验证失败"
assert (
False
), f"未能读取反向边属性: {node2_id} -> {node1_id}，无向图特性验证失败"
print("基本测试完成，数据已保留在数据库中")
return True
@@ -330,7 +334,9 @@ async def test_graph_advanced(storage):
print(f"== 测试反向边的度数: {node2_id} -> {node1_id}")
reverse_edge_degree = await storage.edge_degree(node2_id, node1_id)
print(f"反向边 {node2_id} -> {node1_id} 的度数: {reverse_edge_degree}")
assert edge_degree == reverse_edge_degree, f"正向边和反向边的度数不一致，无向图特性验证失败"
assert (
edge_degree == reverse_edge_degree
), "正向边和反向边的度数不一致，无向图特性验证失败"
print("无向图特性验证成功：正向边和反向边的度数一致")
# 4. 测试 get_node_edges - 获取节点的所有边
@@ -348,14 +354,22 @@ async def test_graph_advanced(storage):
has_connection_with_node3 = False
for edge in node2_edges:
# 检查是否有与node1的连接（无论方向）
if (edge[0] == node1_id and edge[1] == node2_id) or (edge[0] == node2_id and edge[1] == node1_id):
if (edge[0] == node1_id and edge[1] == node2_id) or (
edge[0] == node2_id and edge[1] == node1_id
):
has_connection_with_node1 = True
# 检查是否有与node3的连接（无论方向）
if (edge[0] == node2_id and edge[1] == node3_id) or (edge[0] == node3_id and edge[1] == node2_id):
if (edge[0] == node2_id and edge[1] == node3_id) or (
edge[0] == node3_id and edge[1] == node2_id
):
has_connection_with_node3 = True
assert has_connection_with_node1, f"节点 {node2_id} 的边列表中应包含与 {node1_id} 的连接"
assert has_connection_with_node3, f"节点 {node2_id} 的边列表中应包含与 {node3_id} 的连接"
assert (
has_connection_with_node1
), f"节点 {node2_id} 的边列表中应包含与 {node1_id} 的连接"
assert (
has_connection_with_node3
), f"节点 {node2_id} 的边列表中应包含与 {node3_id} 的连接"
print(f"无向图特性验证成功：节点 {node2_id} 的边列表包含所有相关的边")
# 5. 测试 get_all_labels - 获取所有标签
@@ -398,7 +412,9 @@ async def test_graph_advanced(storage):
print(f"== 验证删除边的无向图特性: {node3_id} -> {node2_id}")
reverse_edge_props = await storage.get_edge(node3_id, node2_id)
print(f"删除后查询反向边属性 {node3_id} -> {node2_id}: {reverse_edge_props}")
assert reverse_edge_props is None, f"反向边 {node3_id} -> {node2_id} 也应被删除，无向图特性验证失败"
assert (
reverse_edge_props is None
), f"反向边 {node3_id} -> {node2_id} 也应被删除，无向图特性验证失败"
print("无向图特性验证成功：删除一个方向的边后，反向边也被删除")
# 9. 测试 remove_nodes - 批量删除节点
@@ -664,12 +680,19 @@ async def test_graph_batch_operations(storage):
reverse_edge_dicts = [{"src": tgt, "tgt": src} for src, tgt in edges]
reverse_edges_dict = await storage.get_edges_batch(reverse_edge_dicts)
print(f"批量获取反向边属性结果: {reverse_edges_dict.keys()}")
assert len(reverse_edges_dict) == 3, f"应返回3条反向边的属性，实际返回 {len(reverse_edges_dict)} 条"
assert (
len(reverse_edges_dict) == 3
), f"应返回3条反向边的属性，实际返回 {len(reverse_edges_dict)} 条"
# 验证正向和反向边的属性是否一致
for (src, tgt), props in edges_dict.items():
assert (tgt, src) in reverse_edges_dict, f"反向边 {tgt} -> {src} 应在返回结果中"
assert props == reverse_edges_dict[(tgt, src)], f"边 {src} -> {tgt} 和反向边 {tgt} -> {src} 的属性不一致"
assert (
tgt,
src,
) in reverse_edges_dict, f"反向边 {tgt} -> {src} 应在返回结果中"
assert (
props == reverse_edges_dict[(tgt, src)]
), f"边 {src} -> {tgt} 和反向边 {tgt} -> {src} 的属性不一致"
print("无向图特性验证成功：批量获取的正向和反向边属性一致")
@@ -693,8 +716,12 @@ async def test_graph_batch_operations(storage):
print("== 验证批量获取节点边的无向图特性")
# 检查节点1的边是否包含所有相关的边（无论方向）
node1_outgoing_edges = [(src, tgt) for src, tgt in nodes_edges[node1_id] if src == node1_id]
node1_incoming_edges = [(src, tgt) for src, tgt in nodes_edges[node1_id] if tgt == node1_id]
node1_outgoing_edges = [
(src, tgt) for src, tgt in nodes_edges[node1_id] if src == node1_id
]
node1_incoming_edges = [
(src, tgt) for src, tgt in nodes_edges[node1_id] if tgt == node1_id
]
print(f"节点 {node1_id} 的出边: {node1_outgoing_edges}")
print(f"节点 {node1_id} 的入边: {node1_incoming_edges}")
@@ -708,8 +735,12 @@ async def test_graph_batch_operations(storage):
assert has_edge_to_node5, f"节点 {node1_id} 的边列表中应包含到 {node5_id} 的边"
# 检查节点3的边是否包含所有相关的边（无论方向）
node3_outgoing_edges = [(src, tgt) for src, tgt in nodes_edges[node3_id] if src == node3_id]
node3_incoming_edges = [(src, tgt) for src, tgt in nodes_edges[node3_id] if tgt == node3_id]
node3_outgoing_edges = [
(src, tgt) for src, tgt in nodes_edges[node3_id] if src == node3_id
]
node3_incoming_edges = [
(src, tgt) for src, tgt in nodes_edges[node3_id] if tgt == node3_id
]
print(f"节点 {node3_id} 的出边: {node3_outgoing_edges}")
print(f"节点 {node3_id} 的入边: {node3_incoming_edges}")
@@ -718,7 +749,9 @@ async def test_graph_batch_operations(storage):
has_edge_to_node4 = any(tgt == node4_id for _, tgt in node3_outgoing_edges)
has_edge_to_node5 = any(tgt == node5_id for _, tgt in node3_outgoing_edges)
assert has_edge_from_node2, f"节点 {node3_id} 的边列表中应包含从 {node2_id} 来的边"
assert (
has_edge_from_node2
), f"节点 {node3_id} 的边列表中应包含从 {node2_id} 来的边"
assert has_edge_to_node4, f"节点 {node3_id} 的边列表中应包含到 {node4_id} 的边"
assert has_edge_to_node5, f"节点 {node3_id} 的边列表中应包含到 {node5_id} 的边"
@@ -810,7 +843,9 @@ async def test_graph_undirected_property(storage):
assert reverse_edge is not None, f"未能读取反向边属性: {node2_id} -> {node1_id}"
# 验证正向和反向边属性是否一致
assert forward_edge == reverse_edge, "正向和反向边属性不一致，无向图特性验证失败"
assert (
forward_edge == reverse_edge
), "正向和反向边属性不一致，无向图特性验证失败"
print("无向图特性验证成功：正向和反向边属性一致")
# 3. 测试边的度数的无向图特性
@@ -830,7 +865,9 @@ async def test_graph_undirected_property(storage):
reverse_degree = await storage.edge_degree(node2_id, node1_id)
print(f"正向边 {node1_id} -> {node2_id} 的度数: {forward_degree}")
print(f"反向边 {node2_id} -> {node1_id} 的度数: {reverse_degree}")
assert forward_degree == reverse_degree, "正向和反向边的度数不一致，无向图特性验证失败"
assert (
forward_degree == reverse_degree
), "正向和反向边的度数不一致，无向图特性验证失败"
print("无向图特性验证成功：正向和反向边的度数一致")
# 4. 测试删除边的无向图特性
@@ -848,7 +885,9 @@ async def test_graph_undirected_property(storage):
# 验证反向边是否也被删除
reverse_edge = await storage.get_edge(node2_id, node1_id)
print(f"删除后查询反向边属性 {node2_id} -> {node1_id}: {reverse_edge}")
assert reverse_edge is None, f"反向边 {node2_id} -> {node1_id} 也应被删除，无向图特性验证失败"
assert (
reverse_edge is None
), f"反向边 {node2_id} -> {node1_id} 也应被删除，无向图特性验证失败"
print("无向图特性验证成功：删除一个方向的边后，反向边也被删除")
# 5. 测试批量操作中的无向图特性
@@ -858,8 +897,14 @@ async def test_graph_undirected_property(storage):
await storage.upsert_edge(node1_id, node2_id, edge1_data)
# 批量获取边属性
edge_dicts = [{"src": node1_id, "tgt": node2_id}, {"src": node1_id, "tgt": node3_id}]
reverse_edge_dicts = [{"src": node2_id, "tgt": node1_id}, {"src": node3_id, "tgt": node1_id}]
edge_dicts = [
{"src": node1_id, "tgt": node2_id},
{"src": node1_id, "tgt": node3_id},
]
reverse_edge_dicts = [
{"src": node2_id, "tgt": node1_id},
{"src": node3_id, "tgt": node1_id},
]
edges_dict = await storage.get_edges_batch(edge_dicts)
reverse_edges_dict = await storage.get_edges_batch(reverse_edge_dicts)
@@ -869,8 +914,13 @@ async def test_graph_undirected_property(storage):
# 验证正向和反向边的属性是否一致
for (src, tgt), props in edges_dict.items():
assert (tgt, src) in reverse_edges_dict, f"反向边 {tgt} -> {src} 应在返回结果中"
assert props == reverse_edges_dict[(tgt, src)], f"边 {src} -> {tgt} 和反向边 {tgt} -> {src} 的属性不一致"
assert (
tgt,
src,
) in reverse_edges_dict, f"反向边 {tgt} -> {src} 应在返回结果中"
assert (
props == reverse_edges_dict[(tgt, src)]
), f"边 {src} -> {tgt} 和反向边 {tgt} -> {src} 的属性不一致"
print("无向图特性验证成功：批量获取的正向和反向边属性一致")
@@ -885,15 +935,25 @@ async def test_graph_undirected_property(storage):
node2_edges = nodes_edges[node2_id]
# 检查节点1是否有到节点2和节点3的边
has_edge_to_node2 = any((src == node1_id and tgt == node2_id) for src, tgt in node1_edges)
has_edge_to_node3 = any((src == node1_id and tgt == node3_id) for src, tgt in node1_edges)
has_edge_to_node2 = any(
(src == node1_id and tgt == node2_id) for src, tgt in node1_edges
)
has_edge_to_node3 = any(
(src == node1_id and tgt == node3_id) for src, tgt in node1_edges
)
assert has_edge_to_node2, f"节点 {node1_id} 的边列表中应包含到 {node2_id} 的边"
assert has_edge_to_node3, f"节点 {node1_id} 的边列表中应包含到 {node3_id} 的边"
# 检查节点2是否有到节点1的边
has_edge_to_node1 = any((src == node2_id and tgt == node1_id) or (src == node1_id and tgt == node2_id) for src, tgt in node2_edges)
assert has_edge_to_node1, f"节点 {node2_id} 的边列表中应包含与 {node1_id} 的连接"
has_edge_to_node1 = any(
(src == node2_id and tgt == node1_id)
or (src == node1_id and tgt == node2_id)
for src, tgt in node2_edges
)
assert (
has_edge_to_node1
), f"节点 {node2_id} 的边列表中应包含与 {node1_id} 的连接"
print("无向图特性验证成功：批量获取的节点边包含所有相关的边（无论方向）")