gzdaniel commited on
Commit
028d831
·
1 Parent(s): 459b6e2

Remove deprecated code from Postgres_impl.py

Browse files

- Stop filtering out 'base' node labels
- Match any edge type in query to improve performance

Files changed (1) hide show
  1. lightrag/kg/postgres_impl.py +7 -6
lightrag/kg/postgres_impl.py CHANGED
@@ -1701,10 +1701,11 @@ class PGGraphStorage(BaseGraphStorage):
1701
  )
1702
 
1703
  # Remove the 'base' label if present in a 'labels' property
1704
- if "labels" in node_dict:
1705
- node_dict["labels"] = [
1706
- label for label in node_dict["labels"] if label != "base"
1707
- ]
 
1708
  nodes_dict[result["node_id"]] = node_dict
1709
 
1710
  return nodes_dict
@@ -1833,14 +1834,14 @@ class PGGraphStorage(BaseGraphStorage):
1833
  forward_query = f"""SELECT * FROM cypher('{self.graph_name}', $$
1834
  WITH [{src_array}] AS sources, [{tgt_array}] AS targets
1835
  UNWIND range(0, size(sources)-1) AS i
1836
- MATCH (a:base {{entity_id: sources[i]}})-[r:DIRECTED]->(b:base {{entity_id: targets[i]}})
1837
  RETURN sources[i] AS source, targets[i] AS target, properties(r) AS edge_properties
1838
  $$) AS (source text, target text, edge_properties agtype)"""
1839
 
1840
  backward_query = f"""SELECT * FROM cypher('{self.graph_name}', $$
1841
  WITH [{src_array}] AS sources, [{tgt_array}] AS targets
1842
  UNWIND range(0, size(sources)-1) AS i
1843
- MATCH (a:base {{entity_id: sources[i]}})<-[r:DIRECTED]-(b:base {{entity_id: targets[i]}})
1844
  RETURN sources[i] AS source, targets[i] AS target, properties(r) AS edge_properties
1845
  $$) AS (source text, target text, edge_properties agtype)"""
1846
 
 
1701
  )
1702
 
1703
  # Remove the 'base' label if present in a 'labels' property
1704
+ # if "labels" in node_dict:
1705
+ # node_dict["labels"] = [
1706
+ # label for label in node_dict["labels"] if label != "base"
1707
+ # ]
1708
+
1709
  nodes_dict[result["node_id"]] = node_dict
1710
 
1711
  return nodes_dict
 
1834
  forward_query = f"""SELECT * FROM cypher('{self.graph_name}', $$
1835
  WITH [{src_array}] AS sources, [{tgt_array}] AS targets
1836
  UNWIND range(0, size(sources)-1) AS i
1837
+ MATCH (a:base {{entity_id: sources[i]}})-[r]->(b:base {{entity_id: targets[i]}})
1838
  RETURN sources[i] AS source, targets[i] AS target, properties(r) AS edge_properties
1839
  $$) AS (source text, target text, edge_properties agtype)"""
1840
 
1841
  backward_query = f"""SELECT * FROM cypher('{self.graph_name}', $$
1842
  WITH [{src_array}] AS sources, [{tgt_array}] AS targets
1843
  UNWIND range(0, size(sources)-1) AS i
1844
+ MATCH (a:base {{entity_id: sources[i]}})<-[r]-(b:base {{entity_id: targets[i]}})
1845
  RETURN sources[i] AS source, targets[i] AS target, properties(r) AS edge_properties
1846
  $$) AS (source text, target text, edge_properties agtype)"""
1847