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

AGEStorage: fix debug logging, and some cosmetics

Browse Source
This commit is contained in:
Alex Potapenko
2024-12-14 12:47:29 +01:00
parent eb2e6d0d87
commit 05c73bcfd1
1 changed files with 64 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

102
lightrag/kg/age_impl.py
View File

@@ -282,11 +282,9 @@ class AGEStorage(BaseGraphStorage):
select_str = "*" select_str = "*"
query = query.format(**params)
return template.format( return template.format(
graph_name=graph_name, graph_name=graph_name,
query=query, query=query.format(**params),
fields=fields_str, fields=fields_str,
projection=select_str, projection=select_str,
) )
@@ -349,16 +347,15 @@ class AGEStorage(BaseGraphStorage):
async def has_node(self, node_id: str) -> bool: async def has_node(self, node_id: str) -> bool:
entity_name_label = node_id.strip('"') entity_name_label = node_id.strip('"')
query = "MATCH (n:`{label}`) RETURN count(n) > 0 AS node_exists" query = """
single_result = ( MATCH (n:`{label}`) RETURN count(n) > 0 AS node_exists
await self._query( """
query, label=AGEStorage._encode_graph_label(entity_name_label) params = {"label": AGEStorage._encode_graph_label(entity_name_label)}
) single_result = (await self._query(query, **params))[0]
)[0]
logger.debug( logger.debug(
"{%s}:query:{%s}:result:{%s}", "{%s}:query:{%s}:result:{%s}",
inspect.currentframe().f_code.co_name, inspect.currentframe().f_code.co_name,
query, query.format(**params),
single_result[0], single_result[0],
) )
@@ -368,20 +365,19 @@ class AGEStorage(BaseGraphStorage):
entity_name_label_source = source_node_id.strip('"') entity_name_label_source = source_node_id.strip('"')
entity_name_label_target = target_node_id.strip('"') entity_name_label_target = target_node_id.strip('"')
query = ( query = """
"MATCH (a:`{src_label}`)-[r]-(b:`{tgt_label}`) " MATCH (a:`{src_label}`)-[r]-(b:`{tgt_label}`)
"RETURN COUNT(r) > 0 AS edge_exists" RETURN COUNT(r) > 0 AS edge_exists
) """
single_result = ( params = {
await self._query( "src_label": AGEStorage._encode_graph_label(entity_name_label_source),
query, "tgt_label": AGEStorage._encode_graph_label(entity_name_label_target),
src_label=AGEStorage._encode_graph_label(entity_name_label_source), }
tgt_label=AGEStorage._encode_graph_label(entity_name_label_target), single_result = (await self._query(query, **params))[0]
)
)[0]
logger.debug( logger.debug(
"{%s}:query:{query}:result:{%s}", "{%s}:query:{%s}:result:{%s}",
inspect.currentframe().f_code.co_name, inspect.currentframe().f_code.co_name,
query.format(**params),
single_result[0], single_result[0],
) )
return single_result["edge_exists"].lower() == "true" return single_result["edge_exists"].lower() == "true"
@@ -389,16 +385,15 @@ class AGEStorage(BaseGraphStorage):
async def get_node(self, node_id: str) -> Union[dict, None]: async def get_node(self, node_id: str) -> Union[dict, None]:
entity_name_label = node_id.strip('"') entity_name_label = node_id.strip('"')
query = "MATCH (n:`{label}`) RETURN n" query = "MATCH (n:`{label}`) RETURN n"
record = await self._query( params = {"label": AGEStorage._encode_graph_label(entity_name_label)}
query, label=AGEStorage._encode_graph_label(entity_name_label) record = await self._query(query, **params)
)
if record: if record:
node = record[0] node = record[0]
node_dict = node["n"] node_dict = node["n"]
logger.debug( logger.debug(
"{%s}: query: {%s}, result: {%s}", "{%s}: query: {%s}, result: {%s}",
inspect.currentframe().f_code.co_name, inspect.currentframe().f_code.co_name,
query, query.format(**params),
node_dict, node_dict,
) )
return node_dict return node_dict
@@ -411,17 +406,14 @@ class AGEStorage(BaseGraphStorage):
MATCH (n:`{label}`)-[]->(x) MATCH (n:`{label}`)-[]->(x)
RETURN count(x) AS total_edge_count RETURN count(x) AS total_edge_count
""" """
record = ( params = {"label": AGEStorage._encode_graph_label(entity_name_label)}
await self._query( record = (await self._query(query, **params))[0]
query, label=AGEStorage._encode_graph_label(entity_name_label)
)
)[0]
if record: if record:
edge_count = int(record["total_edge_count"]) edge_count = int(record["total_edge_count"])
logger.debug( logger.debug(
"{%s}:query:{%s}:result:{%s}", "{%s}:query:{%s}:result:{%s}",
inspect.currentframe().f_code.co_name, inspect.currentframe().f_code.co_name,
query, query.format(**params),
edge_count, edge_count,
) )
return edge_count return edge_count
@@ -465,18 +457,17 @@ class AGEStorage(BaseGraphStorage):
RETURN properties(r) as edge_properties RETURN properties(r) as edge_properties
LIMIT 1 LIMIT 1
""" """
params = {
record = await self._query( "src_label": AGEStorage._encode_graph_label(entity_name_label_source),
query, "tgt_label": AGEStorage._encode_graph_label(entity_name_label_target),
src_label=AGEStorage._encode_graph_label(entity_name_label_source), }
tgt_label=AGEStorage._encode_graph_label(entity_name_label_target), record = await self._query(query, **params)
)
if record and record[0] and record[0]["edge_properties"]: if record and record[0] and record[0]["edge_properties"]:
result = record[0]["edge_properties"] result = record[0]["edge_properties"]
logger.debug( logger.debug(
"{%s}:query:{%s}:result:{%s}", "{%s}:query:{%s}:result:{%s}",
inspect.currentframe().f_code.co_name, inspect.currentframe().f_code.co_name,
query, query.format(**params),
result, result,
) )
return result return result
@@ -488,12 +479,13 @@ class AGEStorage(BaseGraphStorage):
""" """
node_label = source_node_id.strip('"') node_label = source_node_id.strip('"')
query = """MATCH (n:`{label}`) query = """
MATCH (n:`{label}`)
OPTIONAL MATCH (n)-[r]-(connected) OPTIONAL MATCH (n)-[r]-(connected)
RETURN n, r, connected""" RETURN n, r, connected
results = await self._query( """
query, label=AGEStorage._encode_graph_label(node_label) params = {"label": AGEStorage._encode_graph_label(node_label)}
) results = await self._query(query, **params)
edges = [] edges = []
for record in results: for record in results:
source_node = record["n"] if record["n"] else None source_node = record["n"] if record["n"] else None
@@ -533,12 +525,12 @@ class AGEStorage(BaseGraphStorage):
MERGE (n:`{label}`) MERGE (n:`{label}`)
SET n += {properties} SET n += {properties}
""" """
params = {
"label": AGEStorage._encode_graph_label(label),
"properties": AGEStorage._format_properties(properties),
}
try: try:
await self._query( await self._query(query, **params)
query,
label=AGEStorage._encode_graph_label(label),
properties=AGEStorage._format_properties(properties),
)
logger.debug( logger.debug(
"Upserted node with label '{%s}' and properties: {%s}", "Upserted node with label '{%s}' and properties: {%s}",
label, label,
@@ -576,13 +568,13 @@ class AGEStorage(BaseGraphStorage):
SET r += {properties} SET r += {properties}
RETURN r RETURN r
""" """
params = {
"src_label": AGEStorage._encode_graph_label(source_node_label),
"tgt_label": AGEStorage._encode_graph_label(target_node_label),
"properties": AGEStorage._format_properties(edge_properties),
}
try: try:
await self._query( await self._query(query, **params)
query,
src_label=AGEStorage._encode_graph_label(source_node_label),
tgt_label=AGEStorage._encode_graph_label(target_node_label),
properties=AGEStorage._format_properties(edge_properties),
)
logger.debug( logger.debug(
"Upserted edge from '{%s}' to '{%s}' with properties: {%s}", "Upserted edge from '{%s}' to '{%s}' with properties: {%s}",
source_node_label, source_node_label,