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

Refactoring PostgreSQL AGE graph db implementation

Browse Source
This commit is contained in:
yangdx
2025-04-03 15:16:48 +08:00
parent 8878c0e998
commit 33f5629d8a
1 changed files with 120 additions and 211 deletions
Download Patch File Download Diff File Expand all files Collapse all files

331
lightrag/kg/postgres_impl.py
View File

@@ -1064,31 +1064,11 @@ class PGGraphStorage(BaseGraphStorage):
if v.startswith("[") and v.endswith("]"):
if "::vertex" in v:
v = v.replace("::vertex", "")
vertexes = json.loads(v)
dl = []
for vertex in vertexes:
prop = vertex.get("properties")
if not prop:
prop = {}
prop["label"] = PGGraphStorage._decode_graph_label(
prop["node_id"]
)
dl.append(prop)
d[k] = dl
d[k] = json.loads(v)
elif "::edge" in v:
v = v.replace("::edge", "")
edges = json.loads(v)
dl = []
for edge in edges:
dl.append(
(
vertices[edge["start_id"]],
edge["label"],
vertices[edge["end_id"]],
)
)
d[k] = dl
d[k] = json.loads(v)
else:
print("WARNING: unsupported type")
continue
@@ -1097,26 +1077,9 @@ class PGGraphStorage(BaseGraphStorage):
dtype = v.split("::")[-1]
v = v.split("::")[0]
if dtype == "vertex":
vertex = json.loads(v)
field = vertex.get("properties")
if not field:
field = {}
field["label"] = PGGraphStorage._decode_graph_label(
field["node_id"]
)
d[k] = field
# convert edge from id-label->id by replacing id with node information
# we only do this if the vertex was also returned in the query
# this is an attempt to be consistent with neo4j implementation
d[k] = json.loads(v)
elif dtype == "edge":
edge = json.loads(v)
d[k] = (
vertices.get(edge["start_id"], {}),
edge[
"label"
], # we don't use decode_graph_label(), since edge label is always "DIRECTED"
vertices.get(edge["end_id"], {}),
)
d[k] = json.loads(v)
else:
d[k] = (
json.loads(v)
@@ -1152,56 +1115,6 @@ class PGGraphStorage(BaseGraphStorage):
)
return "{" + ", ".join(props) + "}"
@staticmethod
def _encode_graph_label(label: str) -> str:
"""
Since AGE supports only alphanumerical labels, we will encode generic label as HEX string
Args:
label (str): the original label
Returns:
str: the encoded label
"""
return "x" + label.encode().hex()
@staticmethod
def _decode_graph_label(encoded_label: str) -> str:
"""
Since AGE supports only alphanumerical labels, we will encode generic label as HEX string
Args:
encoded_label (str): the encoded label
Returns:
str: the decoded label
"""
return bytes.fromhex(encoded_label.removeprefix("x")).decode()
@staticmethod
def _get_col_name(field: str, idx: int) -> str:
"""
Convert a cypher return field to a pgsql select field
If possible keep the cypher column name, but create a generic name if necessary
Args:
field (str): a return field from a cypher query to be formatted for pgsql
idx (int): the position of the field in the return statement
Returns:
str: the field to be used in the pgsql select statement
"""
# remove white space
field = field.strip()
# if an alias is provided for the field, use it
if " as " in field:
return field.split(" as ")[-1].strip()
# if the return value is an unnamed primitive, give it a generic name
if field.isnumeric() or field in ("true", "false", "null"):
return f"column_{idx}"
# otherwise return the value stripping out some common special chars
return field.replace("(", "_").replace(")", "")
async def _query(
self,
query: str,
@@ -1252,10 +1165,10 @@ class PGGraphStorage(BaseGraphStorage):
return result
async def has_node(self, node_id: str) -> bool:
entity_name_label = self._encode_graph_label(node_id.strip('"'))
entity_name_label = node_id.strip('"')
query = """SELECT * FROM cypher('%s', $$
MATCH (n:base {node_id: "%s"})
MATCH (n:base {entity_id: "%s"})
RETURN count(n) > 0 AS node_exists
$$) AS (node_exists bool)""" % (self.graph_name, entity_name_label)
@@ -1264,11 +1177,11 @@ class PGGraphStorage(BaseGraphStorage):
return single_result["node_exists"]
async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
src_label = self._encode_graph_label(source_node_id.strip('"'))
tgt_label = self._encode_graph_label(target_node_id.strip('"'))
src_label = source_node_id.strip('"')
tgt_label = target_node_id.strip('"')
query = """SELECT * FROM cypher('%s', $$
MATCH (a:base {node_id: "%s"})-[r]-(b:base {node_id: "%s"})
MATCH (a:base {entity_id: "%s"})-[r]-(b:base {entity_id: "%s"})
RETURN COUNT(r) > 0 AS edge_exists
$$) AS (edge_exists bool)""" % (
self.graph_name,
@@ -1281,13 +1194,14 @@ class PGGraphStorage(BaseGraphStorage):
return single_result["edge_exists"]
async def get_node(self, node_id: str) -> dict[str, str] | None:
label = self._encode_graph_label(node_id.strip('"'))
label = node_id.strip('"')
query = """SELECT * FROM cypher('%s', $$
MATCH (n:base {node_id: "%s"})
MATCH (n:base {entity_id: "%s"})
RETURN n
$$) AS (n agtype)""" % (self.graph_name, label)
record = await self._query(query)
if record:
print(f"Record: {record}")
node = record[0]
node_dict = node["n"]
@@ -1295,10 +1209,10 @@ class PGGraphStorage(BaseGraphStorage):
return None
async def node_degree(self, node_id: str) -> int:
label = self._encode_graph_label(node_id.strip('"'))
label = node_id.strip('"')
query = """SELECT * FROM cypher('%s', $$
MATCH (n:base {node_id: "%s"})-[]->(x)
MATCH (n:base {entity_id: "%s"})-[]->(x)
RETURN count(x) AS total_edge_count
$$) AS (total_edge_count integer)""" % (self.graph_name, label)
record = (await self._query(query))[0]
@@ -1322,11 +1236,11 @@ class PGGraphStorage(BaseGraphStorage):
async def get_edge(
self, source_node_id: str, target_node_id: str
) -> dict[str, str] | None:
src_label = self._encode_graph_label(source_node_id.strip('"'))
tgt_label = self._encode_graph_label(target_node_id.strip('"'))
src_label = source_node_id.strip('"')
tgt_label = target_node_id.strip('"')
query = """SELECT * FROM cypher('%s', $$
MATCH (a:base {node_id: "%s"})-[r]->(b:base {node_id: "%s"})
MATCH (a:base {entity_id: "%s"})-[r]->(b:base {entity_id: "%s"})
RETURN properties(r) as edge_properties
LIMIT 1
$$) AS (edge_properties agtype)""" % (
@@ -1336,6 +1250,7 @@ class PGGraphStorage(BaseGraphStorage):
)
record = await self._query(query)
if record and record[0] and record[0]["edge_properties"]:
print(f"Record: {record}")
result = record[0]["edge_properties"]
return result
@@ -1345,10 +1260,10 @@ class PGGraphStorage(BaseGraphStorage):
Retrieves all edges (relationships) for a particular node identified by its label.
:return: list of dictionaries containing edge information
"""
label = self._encode_graph_label(source_node_id.strip('"'))
label = source_node_id.strip('"')
query = """SELECT * FROM cypher('%s', $$
MATCH (n:base {node_id: "%s"})
MATCH (n:base {entity_id: "%s"})
OPTIONAL MATCH (n)-[]-(connected:base)
RETURN n, connected
$$) AS (n agtype, connected agtype)""" % (
@@ -1362,24 +1277,17 @@ class PGGraphStorage(BaseGraphStorage):
source_node = record["n"] if record["n"] else None
connected_node = record["connected"] if record["connected"] else None
source_label = (
source_node["node_id"]
if source_node and source_node["node_id"]
else None
)
target_label = (
connected_node["node_id"]
if connected_node and connected_node["node_id"]
else None
)
if (
source_node
and connected_node
and "properties" in source_node
and "properties" in connected_node
):
source_label = source_node["properties"].get("entity_id")
target_label = connected_node["properties"].get("entity_id")
if source_label and target_label:
edges.append(
(
self._decode_graph_label(source_label),
self._decode_graph_label(target_label),
)
)
if source_label and target_label:
edges.append((source_label, target_label))
return edges
@@ -1389,17 +1297,17 @@ class PGGraphStorage(BaseGraphStorage):
retry=retry_if_exception_type((PGGraphQueryException,)),
)
async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
label = self._encode_graph_label(node_id.strip('"'))
properties = node_data
label = node_id.strip('"')
properties = self._format_properties(node_data)
query = """SELECT * FROM cypher('%s', $$
MERGE (n:base {node_id: "%s"})
MERGE (n:base {entity_id: "%s"})
SET n += %s
RETURN n
$$) AS (n agtype)""" % (
self.graph_name,
label,
self._format_properties(properties),
properties,
)
try:
@@ -1425,14 +1333,14 @@ class PGGraphStorage(BaseGraphStorage):
target_node_id (str): Label of the target node (used as identifier)
edge_data (dict): dictionary of properties to set on the edge
"""
src_label = self._encode_graph_label(source_node_id.strip('"'))
tgt_label = self._encode_graph_label(target_node_id.strip('"'))
edge_properties = edge_data
src_label = source_node_id.strip('"')
tgt_label = target_node_id.strip('"')
edge_properties = self._format_properties(edge_data)
query = """SELECT * FROM cypher('%s', $$
MATCH (source:base {node_id: "%s"})
MATCH (source:base {entity_id: "%s"})
WITH source
MATCH (target:base {node_id: "%s"})
MATCH (target:base {entity_id: "%s"})
MERGE (source)-[r:DIRECTED]->(target)
SET r += %s
RETURN r
@@ -1440,7 +1348,7 @@ class PGGraphStorage(BaseGraphStorage):
self.graph_name,
src_label,
tgt_label,
self._format_properties(edge_properties),
edge_properties,
)
try:
@@ -1460,7 +1368,7 @@ class PGGraphStorage(BaseGraphStorage):
Args:
node_id (str): The ID of the node to delete.
"""
label = self._encode_graph_label(node_id.strip('"'))
label = node_id.strip('"')
query = """SELECT * FROM cypher('%s', $$
MATCH (n:base {entity_id: "%s"})
@@ -1480,14 +1388,12 @@ class PGGraphStorage(BaseGraphStorage):
Args:
node_ids (list[str]): A list of node IDs to remove.
"""
encoded_node_ids = [
self._encode_graph_label(node_id.strip('"')) for node_id in node_ids
]
node_id_list = ", ".join([f'"{node_id}"' for node_id in encoded_node_ids])
node_ids = [node_id.strip('"') for node_id in node_ids]
node_id_list = ", ".join([f'"{node_id}"' for node_id in node_ids])
query = """SELECT * FROM cypher('%s', $$
MATCH (n:base)
WHERE n.nentity_id IN [%s]
WHERE n.entity_id IN [%s]
DETACH DELETE n
$$) AS (n agtype)""" % (self.graph_name, node_id_list)
@@ -1505,11 +1411,11 @@ class PGGraphStorage(BaseGraphStorage):
edges (list[tuple[str, str]]): A list of edges to remove, where each edge is a tuple of (source_node_id, target_node_id).
"""
for source, target in edges:
src_label = self._encode_graph_label(source.strip('"'))
tgt_label = self._encode_graph_label(target.strip('"'))
src_label = source.strip('"')
tgt_label = target.strip('"')
query = """SELECT * FROM cypher('%s', $$
MATCH (a:base {node_id: "%s"})-[r]->(b:base {node_id: "%s"})
MATCH (a:base {entity_id: "%s"})-[r]->(b:base {entity_id: "%s"})
DELETE r
$$) AS (r agtype)""" % (self.graph_name, src_label, tgt_label)
@@ -1560,95 +1466,98 @@ class PGGraphStorage(BaseGraphStorage):
return await embed_func()
async def get_knowledge_graph(
self, node_label: str, max_depth: int = 5
self,
node_label: str,
max_depth: int = 3,
max_nodes: int = MAX_GRAPH_NODES,
) -> KnowledgeGraph:
"""
Retrieve a subgraph containing the specified node and its neighbors up to the specified depth.
Retrieve a connected subgraph of nodes where the label includes the specified `node_label`.
Args:
node_label (str): The label of the node to start from. If "*", the entire graph is returned.
max_depth (int): The maximum depth to traverse from the starting node.
node_label: Label of the starting node, * means all nodes
max_depth: Maximum depth of the subgraph, Defaults to 3
max_nodes: Maxiumu nodes to return by BFS, Defaults to 1000
Returns:
KnowledgeGraph: The retrieved subgraph.
KnowledgeGraph object containing nodes and edges, with an is_truncated flag
indicating whether the graph was truncated due to max_nodes limit
"""
MAX_GRAPH_NODES = 1000
# Build the query based on whether we want the full graph or a specific subgraph.
if node_label == "*":
query = f"""SELECT * FROM cypher('{self.graph_name}', $$
MATCH (n:base)
OPTIONAL MATCH (n)-[r]->(m:base)
RETURN n, r, m
LIMIT {MAX_GRAPH_NODES}
$$) AS (n agtype, r agtype, m agtype)"""
MATCH (n:base)
OPTIONAL MATCH (n)-[r]->(target:base)
RETURN collect(distinct n) AS n, collect(distinct r) AS r
LIMIT {MAX_GRAPH_NODES}
$$) AS (n agtype, r agtype)"""
else:
encoded_label = self._encode_graph_label(node_label.strip('"'))
strip_label = node_label.strip('"')
query = f"""SELECT * FROM cypher('{self.graph_name}', $$
MATCH (n:base {{entity_id: "{encoded_label}"}})
OPTIONAL MATCH p = (n)-[*..{max_depth}]-(m)
RETURN nodes(p) AS nodes, relationships(p) AS relationships
LIMIT {MAX_GRAPH_NODES}
$$) AS (nodes agtype, relationships agtype)"""
MATCH (n:base {{entity_id: "{strip_label}"}})
OPTIONAL MATCH p = (n)-[*..{max_depth}]-(m)
RETURN nodes(p) AS n, relationships(p) AS r
LIMIT {max_nodes}
$$) AS (n agtype, r agtype)"""
results = await self._query(query)
nodes = {}
edges = []
unique_edge_ids = set()
def add_node(node_data: dict):
node_id = self._decode_graph_label(node_data["node_id"])
if node_id not in nodes:
nodes[node_id] = node_data
def add_edge(edge_data: list):
src_id = self._decode_graph_label(edge_data[0]["node_id"])
tgt_id = self._decode_graph_label(edge_data[2]["node_id"])
edge_key = f"{src_id},{tgt_id}"
if edge_key not in unique_edge_ids:
unique_edge_ids.add(edge_key)
edges.append(
(
edge_key,
src_id,
tgt_id,
{"source": edge_data[0], "target": edge_data[2]},
# Process the query results with deduplication by node and edge IDs
nodes_dict = {}
edges_dict = {}
for result in results:
# Handle single node cases
if result.get("n") and isinstance(result["n"], dict):
node_id = str(result["n"]["id"])
if node_id not in nodes_dict:
nodes_dict[node_id] = KnowledgeGraphNode(
id=node_id,
labels=[result["n"]["properties"]["entity_id"]],
properties=result["n"]["properties"],
)
)
# Handle node list cases
elif result.get("n") and isinstance(result["n"], list):
for node in result["n"]:
if isinstance(node, dict) and "id" in node:
node_id = str(node["id"])
if node_id not in nodes_dict and "properties" in node:
nodes_dict[node_id] = KnowledgeGraphNode(
id=node_id,
labels=[node["properties"]["entity_id"]],
properties=node["properties"],
)
# Handle single edge cases
if result.get("r") and isinstance(result["r"], dict):
edge_id = str(result["r"]["id"])
if edge_id not in edges_dict:
edges_dict[edge_id] = KnowledgeGraphEdge(
id=edge_id,
type="DIRECTED",
source=str(result["r"]["start_id"]),
target=str(result["r"]["end_id"]),
properties=result["r"]["properties"],
)
# Handle edge list cases
elif result.get("r") and isinstance(result["r"], list):
for edge in result["r"]:
if isinstance(edge, dict) and "id" in edge:
edge_id = str(edge["id"])
if edge_id not in edges_dict:
edges_dict[edge_id] = KnowledgeGraphEdge(
id=edge_id,
type="DIRECTED",
source=str(edge["start_id"]),
target=str(edge["end_id"]),
properties=edge["properties"],
)
# Process the query results.
if node_label == "*":
for result in results:
if result.get("n"):
add_node(result["n"])
if result.get("m"):
add_node(result["m"])
if result.get("r"):
add_edge(result["r"])
else:
for result in results:
for node in result.get("nodes", []):
add_node(node)
for edge in result.get("relationships", []):
add_edge(edge)
# Construct and return the KnowledgeGraph.
# Construct and return the KnowledgeGraph with deduplicated nodes and edges
kg = KnowledgeGraph(
nodes=[
KnowledgeGraphNode(id=node_id, labels=[node_id], properties=node_data)
for node_id, node_data in nodes.items()
],
edges=[
KnowledgeGraphEdge(
id=edge_id,
type="DIRECTED",
source=src,
target=tgt,
properties=props,
)
for edge_id, src, tgt, props in edges
],
nodes=list(nodes_dict.values()),
edges=list(edges_dict.values()),
is_truncated=False,
)
return kg