Remove deprecated storage
This commit is contained in:
yangdx
parent
c22315ea6d
commit
a04c11a598
4 changed files with 0 additions and 2796 deletions
|
|
@ -5,7 +5,6 @@ STORAGE_IMPLEMENTATIONS = {
|
|||
"RedisKVStorage",
|
||||
"PGKVStorage",
|
||||
"MongoKVStorage",
|
||||
# "TiDBKVStorage",
|
||||
],
|
||||
"required_methods": ["get_by_id", "upsert"],
|
||||
},
|
||||
|
|
@ -16,9 +15,6 @@ STORAGE_IMPLEMENTATIONS = {
|
|||
"PGGraphStorage",
|
||||
"MongoGraphStorage",
|
||||
"MemgraphStorage",
|
||||
# "AGEStorage",
|
||||
# "TiDBGraphStorage",
|
||||
# "GremlinStorage",
|
||||
],
|
||||
"required_methods": ["upsert_node", "upsert_edge"],
|
||||
},
|
||||
|
|
@ -31,7 +27,6 @@ STORAGE_IMPLEMENTATIONS = {
|
|||
"QdrantVectorDBStorage",
|
||||
"MongoVectorDBStorage",
|
||||
# "ChromaVectorDBStorage",
|
||||
# "TiDBVectorDBStorage",
|
||||
],
|
||||
"required_methods": ["query", "upsert"],
|
||||
},
|
||||
|
|
@ -52,20 +47,17 @@ STORAGE_ENV_REQUIREMENTS: dict[str, list[str]] = {
|
|||
"JsonKVStorage": [],
|
||||
"MongoKVStorage": [],
|
||||
"RedisKVStorage": ["REDIS_URI"],
|
||||
# "TiDBKVStorage": ["TIDB_USER", "TIDB_PASSWORD", "TIDB_DATABASE"],
|
||||
"PGKVStorage": ["POSTGRES_USER", "POSTGRES_PASSWORD", "POSTGRES_DATABASE"],
|
||||
# Graph Storage Implementations
|
||||
"NetworkXStorage": [],
|
||||
"Neo4JStorage": ["NEO4J_URI", "NEO4J_USERNAME", "NEO4J_PASSWORD"],
|
||||
"MongoGraphStorage": [],
|
||||
"MemgraphStorage": ["MEMGRAPH_URI"],
|
||||
# "TiDBGraphStorage": ["TIDB_USER", "TIDB_PASSWORD", "TIDB_DATABASE"],
|
||||
"AGEStorage": [
|
||||
"AGE_POSTGRES_DB",
|
||||
"AGE_POSTGRES_USER",
|
||||
"AGE_POSTGRES_PASSWORD",
|
||||
],
|
||||
# "GremlinStorage": ["GREMLIN_HOST", "GREMLIN_PORT", "GREMLIN_GRAPH"],
|
||||
"PGGraphStorage": [
|
||||
"POSTGRES_USER",
|
||||
"POSTGRES_PASSWORD",
|
||||
|
|
@ -75,7 +67,6 @@ STORAGE_ENV_REQUIREMENTS: dict[str, list[str]] = {
|
|||
"NanoVectorDBStorage": [],
|
||||
"MilvusVectorDBStorage": [],
|
||||
"ChromaVectorDBStorage": [],
|
||||
# "TiDBVectorDBStorage": ["TIDB_USER", "TIDB_PASSWORD", "TIDB_DATABASE"],
|
||||
"PGVectorStorage": ["POSTGRES_USER", "POSTGRES_PASSWORD", "POSTGRES_DATABASE"],
|
||||
"FaissVectorDBStorage": [],
|
||||
"QdrantVectorDBStorage": ["QDRANT_URL"], # QDRANT_API_KEY has default value None
|
||||
|
|
@ -102,14 +93,10 @@ STORAGES = {
|
|||
"RedisKVStorage": ".kg.redis_impl",
|
||||
"RedisDocStatusStorage": ".kg.redis_impl",
|
||||
"ChromaVectorDBStorage": ".kg.chroma_impl",
|
||||
# "TiDBKVStorage": ".kg.tidb_impl",
|
||||
# "TiDBVectorDBStorage": ".kg.tidb_impl",
|
||||
# "TiDBGraphStorage": ".kg.tidb_impl",
|
||||
"PGKVStorage": ".kg.postgres_impl",
|
||||
"PGVectorStorage": ".kg.postgres_impl",
|
||||
"AGEStorage": ".kg.age_impl",
|
||||
"PGGraphStorage": ".kg.postgres_impl",
|
||||
# "GremlinStorage": ".kg.gremlin_impl",
|
||||
"PGDocStatusStorage": ".kg.postgres_impl",
|
||||
"FaissVectorDBStorage": ".kg.faiss_impl",
|
||||
"QdrantVectorDBStorage": ".kg.qdrant_impl",
|
||||
|
|
|
|||
|
|
@ -1,867 +0,0 @@
|
|||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, NamedTuple, Optional, Union, final
|
||||
import pipmaster as pm
|
||||
from lightrag.types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
|
||||
|
||||
from tenacity import (
|
||||
retry,
|
||||
retry_if_exception_type,
|
||||
stop_after_attempt,
|
||||
wait_exponential,
|
||||
)
|
||||
|
||||
from lightrag.utils import logger
|
||||
|
||||
from ..base import BaseGraphStorage
|
||||
|
||||
if sys.platform.startswith("win"):
|
||||
import asyncio.windows_events
|
||||
|
||||
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
|
||||
|
||||
|
||||
if not pm.is_installed("psycopg-pool"):
|
||||
pm.install("psycopg-pool")
|
||||
pm.install("psycopg[binary,pool]")
|
||||
|
||||
if not pm.is_installed("asyncpg"):
|
||||
pm.install("asyncpg")
|
||||
|
||||
import psycopg # type: ignore
|
||||
from psycopg.rows import namedtuple_row # type: ignore
|
||||
from psycopg_pool import AsyncConnectionPool, PoolTimeout # type: ignore
|
||||
|
||||
|
||||
class AGEQueryException(Exception):
|
||||
"""Exception for the AGE queries."""
|
||||
|
||||
def __init__(self, exception: Union[str, Dict]) -> None:
|
||||
if isinstance(exception, dict):
|
||||
self.message = exception["message"] if "message" in exception else "unknown"
|
||||
self.details = exception["details"] if "details" in exception else "unknown"
|
||||
else:
|
||||
self.message = exception
|
||||
self.details = "unknown"
|
||||
|
||||
def get_message(self) -> str:
|
||||
return self.message
|
||||
|
||||
def get_details(self) -> Any:
|
||||
return self.details
|
||||
|
||||
|
||||
@final
|
||||
@dataclass
|
||||
class AGEStorage(BaseGraphStorage):
|
||||
@staticmethod
|
||||
def load_nx_graph(file_name):
|
||||
print("no preloading of graph with AGE in production")
|
||||
|
||||
def __init__(self, namespace, global_config, embedding_func):
|
||||
super().__init__(
|
||||
namespace=namespace,
|
||||
global_config=global_config,
|
||||
embedding_func=embedding_func,
|
||||
)
|
||||
self._driver = None
|
||||
self._driver_lock = asyncio.Lock()
|
||||
DB = os.environ["AGE_POSTGRES_DB"].replace("\\", "\\\\").replace("'", "\\'")
|
||||
USER = os.environ["AGE_POSTGRES_USER"].replace("\\", "\\\\").replace("'", "\\'")
|
||||
PASSWORD = (
|
||||
os.environ["AGE_POSTGRES_PASSWORD"]
|
||||
.replace("\\", "\\\\")
|
||||
.replace("'", "\\'")
|
||||
)
|
||||
HOST = os.environ["AGE_POSTGRES_HOST"].replace("\\", "\\\\").replace("'", "\\'")
|
||||
PORT = os.environ.get("AGE_POSTGRES_PORT", "8529")
|
||||
self.graph_name = namespace or os.environ.get("AGE_GRAPH_NAME", "lightrag")
|
||||
|
||||
connection_string = f"dbname='{DB}' user='{USER}' password='{PASSWORD}' host='{HOST}' port={PORT}"
|
||||
|
||||
self._driver = AsyncConnectionPool(connection_string, open=False)
|
||||
|
||||
return None
|
||||
|
||||
async def close(self):
|
||||
if self._driver:
|
||||
await self._driver.close()
|
||||
self._driver = None
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
if self._driver:
|
||||
await self._driver.close()
|
||||
|
||||
@staticmethod
|
||||
def _record_to_dict(record: NamedTuple) -> Dict[str, Any]:
|
||||
"""
|
||||
Convert a record returned from an age query to a dictionary
|
||||
|
||||
Args:
|
||||
record (): a record from an age query result
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: a dictionary representation of the record where
|
||||
the dictionary key is the field name and the value is the
|
||||
value converted to a python type
|
||||
"""
|
||||
# result holder
|
||||
d = {}
|
||||
|
||||
# prebuild a mapping of vertex_id to vertex mappings to be used
|
||||
# later to build edges
|
||||
vertices = {}
|
||||
for k in record._fields:
|
||||
v = getattr(record, k)
|
||||
# agtype comes back '{key: value}::type' which must be parsed
|
||||
if isinstance(v, str) and "::" in v:
|
||||
dtype = v.split("::")[-1]
|
||||
v = v.split("::")[0]
|
||||
if dtype == "vertex":
|
||||
vertex = json.loads(v)
|
||||
vertices[vertex["id"]] = vertex.get("properties")
|
||||
|
||||
# iterate returned fields and parse appropriately
|
||||
for k in record._fields:
|
||||
v = getattr(record, k)
|
||||
if isinstance(v, str) and "::" in v:
|
||||
dtype = v.split("::")[-1]
|
||||
v = v.split("::")[0]
|
||||
else:
|
||||
dtype = ""
|
||||
|
||||
if dtype == "vertex":
|
||||
vertex = json.loads(v)
|
||||
field = json.loads(v).get("properties")
|
||||
if not field:
|
||||
field = {}
|
||||
field["label"] = AGEStorage._decode_graph_label(vertex["label"])
|
||||
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
|
||||
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"], {}),
|
||||
)
|
||||
else:
|
||||
d[k] = json.loads(v) if isinstance(v, str) else v
|
||||
|
||||
return d
|
||||
|
||||
@staticmethod
|
||||
def _format_properties(
|
||||
properties: Dict[str, Any], _id: Union[str, None] = None
|
||||
) -> str:
|
||||
"""
|
||||
Convert a dictionary of properties to a string representation that
|
||||
can be used in a cypher query insert/merge statement.
|
||||
|
||||
Args:
|
||||
properties (Dict[str,str]): a dictionary containing node/edge properties
|
||||
id (Union[str, None]): the id of the node or None if none exists
|
||||
|
||||
Returns:
|
||||
str: the properties dictionary as a properly formatted string
|
||||
"""
|
||||
props = []
|
||||
# wrap property key in backticks to escape
|
||||
for k, v in properties.items():
|
||||
prop = f"`{k}`: {json.dumps(v)}"
|
||||
props.append(prop)
|
||||
if _id is not None and "id" not in properties:
|
||||
props.append(
|
||||
f"id: {json.dumps(_id)}" if isinstance(_id, str) else f"id: {_id}"
|
||||
)
|
||||
return "{" + ", ".join(props) + "}"
|
||||
|
||||
@staticmethod
|
||||
def _encode_graph_label(label: str) -> str:
|
||||
"""
|
||||
Since AGE suports 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 suports 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(")", "")
|
||||
|
||||
@staticmethod
|
||||
def _wrap_query(query: str, graph_name: str, **params: str) -> str:
|
||||
"""
|
||||
Convert a cypher query to an Apache Age compatible
|
||||
sql query by wrapping the cypher query in ag_catalog.cypher,
|
||||
casting results to agtype and building a select statement
|
||||
|
||||
Args:
|
||||
query (str): a valid cypher query
|
||||
graph_name (str): the name of the graph to query
|
||||
params (dict): parameters for the query
|
||||
|
||||
Returns:
|
||||
str: an equivalent pgsql query
|
||||
"""
|
||||
|
||||
# pgsql template
|
||||
template = """SELECT {projection} FROM ag_catalog.cypher('{graph_name}', $$
|
||||
{query}
|
||||
$$) AS ({fields});"""
|
||||
|
||||
# if there are any returned fields they must be added to the pgsql query
|
||||
if "return" in query.lower():
|
||||
# parse return statement to identify returned fields
|
||||
fields = (
|
||||
query.lower()
|
||||
.split("return")[-1]
|
||||
.split("distinct")[-1]
|
||||
.split("order by")[0]
|
||||
.split("skip")[0]
|
||||
.split("limit")[0]
|
||||
.split(",")
|
||||
)
|
||||
|
||||
# raise exception if RETURN * is found as we can't resolve the fields
|
||||
if "*" in [x.strip() for x in fields]:
|
||||
raise ValueError(
|
||||
"AGE graph does not support 'RETURN *'"
|
||||
+ " statements in Cypher queries"
|
||||
)
|
||||
|
||||
# get pgsql formatted field names
|
||||
fields = [
|
||||
AGEStorage._get_col_name(field, idx) for idx, field in enumerate(fields)
|
||||
]
|
||||
|
||||
# build resulting pgsql relation
|
||||
fields_str = ", ".join(
|
||||
[field.split(".")[-1] + " agtype" for field in fields]
|
||||
)
|
||||
|
||||
# if no return statement we still need to return a single field of type agtype
|
||||
else:
|
||||
fields_str = "a agtype"
|
||||
|
||||
select_str = "*"
|
||||
|
||||
return template.format(
|
||||
graph_name=graph_name,
|
||||
query=query.format(**params),
|
||||
fields=fields_str,
|
||||
projection=select_str,
|
||||
)
|
||||
|
||||
async def _query(self, query: str, **params: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Query the graph by taking a cypher query, converting it to an
|
||||
age compatible query, executing it and converting the result
|
||||
|
||||
Args:
|
||||
query (str): a cypher query to be executed
|
||||
params (dict): parameters for the query
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: a list of dictionaries containing the result set
|
||||
"""
|
||||
# convert cypher query to pgsql/age query
|
||||
wrapped_query = self._wrap_query(query, self.graph_name, **params)
|
||||
|
||||
await self._driver.open()
|
||||
|
||||
# create graph if it doesn't exist
|
||||
async with self._get_pool_connection() as conn:
|
||||
async with conn.cursor() as curs:
|
||||
try:
|
||||
await curs.execute('SET search_path = ag_catalog, "$user", public')
|
||||
await curs.execute(f"SELECT create_graph('{self.graph_name}')")
|
||||
await conn.commit()
|
||||
except (
|
||||
psycopg.errors.InvalidSchemaName,
|
||||
psycopg.errors.UniqueViolation,
|
||||
):
|
||||
await conn.rollback()
|
||||
|
||||
# execute the query, rolling back on an error
|
||||
async with self._get_pool_connection() as conn:
|
||||
async with conn.cursor(row_factory=namedtuple_row) as curs:
|
||||
try:
|
||||
await curs.execute('SET search_path = ag_catalog, "$user", public')
|
||||
await curs.execute(wrapped_query)
|
||||
await conn.commit()
|
||||
except psycopg.Error as e:
|
||||
await conn.rollback()
|
||||
raise AGEQueryException(
|
||||
{
|
||||
"message": f"Error executing graph query: {query.format(**params)}",
|
||||
"detail": str(e),
|
||||
}
|
||||
) from e
|
||||
|
||||
data = await curs.fetchall()
|
||||
if data is None:
|
||||
result = []
|
||||
# decode records
|
||||
else:
|
||||
result = [AGEStorage._record_to_dict(d) for d in data]
|
||||
|
||||
return result
|
||||
|
||||
async def has_node(self, node_id: str) -> bool:
|
||||
entity_name_label = node_id.strip('"')
|
||||
|
||||
query = """
|
||||
MATCH (n:`{label}`) RETURN count(n) > 0 AS node_exists
|
||||
"""
|
||||
params = {"label": AGEStorage._encode_graph_label(entity_name_label)}
|
||||
single_result = (await self._query(query, **params))[0]
|
||||
logger.debug(
|
||||
"{%s}:query:{%s}:result:{%s}",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
query.format(**params),
|
||||
single_result["node_exists"],
|
||||
)
|
||||
|
||||
return single_result["node_exists"]
|
||||
|
||||
async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
|
||||
entity_name_label_source = source_node_id.strip('"')
|
||||
entity_name_label_target = target_node_id.strip('"')
|
||||
|
||||
query = """
|
||||
MATCH (a:`{src_label}`)-[r]-(b:`{tgt_label}`)
|
||||
RETURN COUNT(r) > 0 AS edge_exists
|
||||
"""
|
||||
params = {
|
||||
"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]
|
||||
logger.debug(
|
||||
"{%s}:query:{%s}:result:{%s}",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
query.format(**params),
|
||||
single_result["edge_exists"],
|
||||
)
|
||||
return single_result["edge_exists"]
|
||||
|
||||
async def get_node(self, node_id: str) -> dict[str, str] | None:
|
||||
entity_name_label = node_id.strip('"')
|
||||
query = """
|
||||
MATCH (n:`{label}`) RETURN n
|
||||
"""
|
||||
params = {"label": AGEStorage._encode_graph_label(entity_name_label)}
|
||||
record = await self._query(query, **params)
|
||||
if record:
|
||||
node = record[0]
|
||||
node_dict = node["n"]
|
||||
logger.debug(
|
||||
"{%s}: query: {%s}, result: {%s}",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
query.format(**params),
|
||||
node_dict,
|
||||
)
|
||||
return node_dict
|
||||
return None
|
||||
|
||||
async def node_degree(self, node_id: str) -> int:
|
||||
entity_name_label = node_id.strip('"')
|
||||
|
||||
query = """
|
||||
MATCH (n:`{label}`)-[]->(x)
|
||||
RETURN count(x) AS total_edge_count
|
||||
"""
|
||||
params = {"label": AGEStorage._encode_graph_label(entity_name_label)}
|
||||
record = (await self._query(query, **params))[0]
|
||||
if record:
|
||||
edge_count = int(record["total_edge_count"])
|
||||
logger.debug(
|
||||
"{%s}:query:{%s}:result:{%s}",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
query.format(**params),
|
||||
edge_count,
|
||||
)
|
||||
return edge_count
|
||||
|
||||
async def edge_degree(self, src_id: str, tgt_id: str) -> int:
|
||||
entity_name_label_source = src_id.strip('"')
|
||||
entity_name_label_target = tgt_id.strip('"')
|
||||
src_degree = await self.node_degree(entity_name_label_source)
|
||||
trg_degree = await self.node_degree(entity_name_label_target)
|
||||
|
||||
# Convert None to 0 for addition
|
||||
src_degree = 0 if src_degree is None else src_degree
|
||||
trg_degree = 0 if trg_degree is None else trg_degree
|
||||
|
||||
degrees = int(src_degree) + int(trg_degree)
|
||||
logger.debug(
|
||||
"{%s}:query:src_Degree+trg_degree:result:{%s}",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
degrees,
|
||||
)
|
||||
return degrees
|
||||
|
||||
async def get_edge(
|
||||
self, source_node_id: str, target_node_id: str
|
||||
) -> dict[str, str] | None:
|
||||
entity_name_label_source = source_node_id.strip('"')
|
||||
entity_name_label_target = target_node_id.strip('"')
|
||||
|
||||
query = """
|
||||
MATCH (a:`{src_label}`)-[r]->(b:`{tgt_label}`)
|
||||
RETURN properties(r) as edge_properties
|
||||
LIMIT 1
|
||||
"""
|
||||
params = {
|
||||
"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"]:
|
||||
result = record[0]["edge_properties"]
|
||||
logger.debug(
|
||||
"{%s}:query:{%s}:result:{%s}",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
query.format(**params),
|
||||
result,
|
||||
)
|
||||
return result
|
||||
|
||||
async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None:
|
||||
"""
|
||||
Retrieves all edges (relationships) for a particular node identified by its label.
|
||||
:return: List of dictionaries containing edge information
|
||||
"""
|
||||
node_label = source_node_id.strip('"')
|
||||
|
||||
query = """
|
||||
MATCH (n:`{label}`)
|
||||
OPTIONAL MATCH (n)-[r]-(connected)
|
||||
RETURN n, r, connected
|
||||
"""
|
||||
params = {"label": AGEStorage._encode_graph_label(node_label)}
|
||||
results = await self._query(query, **params)
|
||||
edges = []
|
||||
for record in results:
|
||||
source_node = record["n"] if record["n"] else None
|
||||
connected_node = record["connected"] if record["connected"] else None
|
||||
|
||||
source_label = (
|
||||
source_node["label"] if source_node and source_node["label"] else None
|
||||
)
|
||||
target_label = (
|
||||
connected_node["label"]
|
||||
if connected_node and connected_node["label"]
|
||||
else None
|
||||
)
|
||||
|
||||
if source_label and target_label:
|
||||
edges.append((source_label, target_label))
|
||||
|
||||
return edges
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential(multiplier=1, min=4, max=10),
|
||||
retry=retry_if_exception_type((AGEQueryException,)),
|
||||
)
|
||||
async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
|
||||
"""
|
||||
Upsert a node in the AGE database.
|
||||
|
||||
Args:
|
||||
node_id: The unique identifier for the node (used as label)
|
||||
node_data: Dictionary of node properties
|
||||
"""
|
||||
label = node_id.strip('"')
|
||||
properties = node_data
|
||||
|
||||
query = """
|
||||
MERGE (n:`{label}`)
|
||||
SET n += {properties}
|
||||
"""
|
||||
params = {
|
||||
"label": AGEStorage._encode_graph_label(label),
|
||||
"properties": AGEStorage._format_properties(properties),
|
||||
}
|
||||
try:
|
||||
await self._query(query, **params)
|
||||
logger.debug(
|
||||
"Upserted node with label '{%s}' and properties: {%s}",
|
||||
label,
|
||||
properties,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Error during upsert: {%s}", e)
|
||||
raise
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential(multiplier=1, min=4, max=10),
|
||||
retry=retry_if_exception_type((AGEQueryException,)),
|
||||
)
|
||||
async def upsert_edge(
|
||||
self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
|
||||
) -> None:
|
||||
"""
|
||||
Upsert an edge and its properties between two nodes identified by their labels.
|
||||
|
||||
Args:
|
||||
source_node_id (str): Label of the source node (used as identifier)
|
||||
target_node_id (str): Label of the target node (used as identifier)
|
||||
edge_data (dict): Dictionary of properties to set on the edge
|
||||
"""
|
||||
source_node_label = source_node_id.strip('"')
|
||||
target_node_label = target_node_id.strip('"')
|
||||
edge_properties = edge_data
|
||||
|
||||
query = """
|
||||
MATCH (source:`{src_label}`)
|
||||
WITH source
|
||||
MATCH (target:`{tgt_label}`)
|
||||
MERGE (source)-[r:DIRECTED]->(target)
|
||||
SET r += {properties}
|
||||
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:
|
||||
await self._query(query, **params)
|
||||
logger.debug(
|
||||
"Upserted edge from '{%s}' to '{%s}' with properties: {%s}",
|
||||
source_node_label,
|
||||
target_node_label,
|
||||
edge_properties,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Error during edge upsert: {%s}", e)
|
||||
raise
|
||||
|
||||
@asynccontextmanager
|
||||
async def _get_pool_connection(self, timeout: Optional[float] = None):
|
||||
"""Workaround for a psycopg_pool bug"""
|
||||
|
||||
try:
|
||||
connection = await self._driver.getconn(timeout=timeout)
|
||||
except PoolTimeout:
|
||||
await self._driver._add_connection(None) # workaround...
|
||||
connection = await self._driver.getconn(timeout=timeout)
|
||||
|
||||
try:
|
||||
async with connection:
|
||||
yield connection
|
||||
finally:
|
||||
await self._driver.putconn(connection)
|
||||
|
||||
async def delete_node(self, node_id: str) -> None:
|
||||
"""Delete a node with the specified label
|
||||
|
||||
Args:
|
||||
node_id: The label of the node to delete
|
||||
"""
|
||||
entity_name_label = node_id.strip('"')
|
||||
|
||||
query = """
|
||||
MATCH (n:`{label}`)
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"label": AGEStorage._encode_graph_label(entity_name_label)}
|
||||
try:
|
||||
await self._query(query, **params)
|
||||
logger.debug(f"Deleted node with label '{entity_name_label}'")
|
||||
except Exception as e:
|
||||
logger.error(f"Error during node deletion: {str(e)}")
|
||||
raise
|
||||
|
||||
async def remove_nodes(self, nodes: list[str]):
|
||||
"""Delete multiple nodes
|
||||
|
||||
Args:
|
||||
nodes: List of node labels to be deleted
|
||||
"""
|
||||
for node in nodes:
|
||||
await self.delete_node(node)
|
||||
|
||||
async def remove_edges(self, edges: list[tuple[str, str]]):
|
||||
"""Delete multiple edges
|
||||
|
||||
Args:
|
||||
edges: List of edges to be deleted, each edge is a (source, target) tuple
|
||||
"""
|
||||
for source, target in edges:
|
||||
entity_name_label_source = source.strip('"')
|
||||
entity_name_label_target = target.strip('"')
|
||||
|
||||
query = """
|
||||
MATCH (source:`{src_label}`)-[r]->(target:`{tgt_label}`)
|
||||
DELETE r
|
||||
"""
|
||||
params = {
|
||||
"src_label": AGEStorage._encode_graph_label(entity_name_label_source),
|
||||
"tgt_label": AGEStorage._encode_graph_label(entity_name_label_target),
|
||||
}
|
||||
try:
|
||||
await self._query(query, **params)
|
||||
logger.debug(
|
||||
f"Deleted edge from '{entity_name_label_source}' to '{entity_name_label_target}'"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during edge deletion: {str(e)}")
|
||||
raise
|
||||
|
||||
async def get_all_labels(self) -> list[str]:
|
||||
"""Get all node labels in the database
|
||||
|
||||
Returns:
|
||||
["label1", "label2", ...] # Alphabetically sorted label list
|
||||
"""
|
||||
query = """
|
||||
MATCH (n)
|
||||
RETURN DISTINCT labels(n) AS node_labels
|
||||
"""
|
||||
results = await self._query(query)
|
||||
|
||||
all_labels = []
|
||||
for record in results:
|
||||
if record and "node_labels" in record:
|
||||
for label in record["node_labels"]:
|
||||
if label:
|
||||
# Decode label
|
||||
decoded_label = AGEStorage._decode_graph_label(label)
|
||||
all_labels.append(decoded_label)
|
||||
|
||||
# Remove duplicates and sort
|
||||
return sorted(list(set(all_labels)))
|
||||
|
||||
async def get_knowledge_graph(
|
||||
self, node_label: str, max_depth: int = 5
|
||||
) -> KnowledgeGraph:
|
||||
"""
|
||||
Retrieve a connected subgraph of nodes where the label includes the specified 'node_label'.
|
||||
Maximum number of nodes is constrained by the environment variable 'MAX_GRAPH_NODES' (default: 1000).
|
||||
When reducing the number of nodes, the prioritization criteria are as follows:
|
||||
1. Label matching nodes take precedence (nodes containing the specified label string)
|
||||
2. Followed by nodes directly connected to the matching nodes
|
||||
3. Finally, the degree of the nodes
|
||||
|
||||
Args:
|
||||
node_label: String to match in node labels (will match any node containing this string in its label)
|
||||
max_depth: Maximum depth of the graph. Defaults to 5.
|
||||
|
||||
Returns:
|
||||
KnowledgeGraph: Complete connected subgraph for specified node
|
||||
"""
|
||||
max_graph_nodes = int(os.getenv("MAX_GRAPH_NODES", 1000))
|
||||
result = KnowledgeGraph()
|
||||
seen_nodes = set()
|
||||
seen_edges = set()
|
||||
|
||||
# Handle special case for "*" label
|
||||
if node_label == "*":
|
||||
# Query all nodes and sort by degree
|
||||
query = """
|
||||
MATCH (n)
|
||||
OPTIONAL MATCH (n)-[r]-()
|
||||
WITH n, count(r) AS degree
|
||||
ORDER BY degree DESC
|
||||
LIMIT {max_nodes}
|
||||
RETURN n, degree
|
||||
"""
|
||||
params = {"max_nodes": max_graph_nodes}
|
||||
nodes_result = await self._query(query, **params)
|
||||
|
||||
# Add nodes to result
|
||||
node_ids = []
|
||||
for record in nodes_result:
|
||||
if "n" in record:
|
||||
node = record["n"]
|
||||
node_id = str(node.get("id", ""))
|
||||
if node_id not in seen_nodes:
|
||||
node_properties = {k: v for k, v in node.items()}
|
||||
node_label = node.get("label", "")
|
||||
result.nodes.append(
|
||||
KnowledgeGraphNode(
|
||||
id=node_id,
|
||||
labels=[node_label],
|
||||
properties=node_properties,
|
||||
)
|
||||
)
|
||||
seen_nodes.add(node_id)
|
||||
node_ids.append(node_id)
|
||||
|
||||
# Query edges between these nodes
|
||||
if node_ids:
|
||||
edges_query = """
|
||||
MATCH (a)-[r]->(b)
|
||||
WHERE a.id IN {node_ids} AND b.id IN {node_ids}
|
||||
RETURN a, r, b
|
||||
"""
|
||||
edges_params = {"node_ids": node_ids}
|
||||
edges_result = await self._query(edges_query, **edges_params)
|
||||
|
||||
# Add edges to result
|
||||
for record in edges_result:
|
||||
if "r" in record and "a" in record and "b" in record:
|
||||
source = record["a"].get("id", "")
|
||||
target = record["b"].get("id", "")
|
||||
edge_id = f"{source}-{target}"
|
||||
if edge_id not in seen_edges:
|
||||
edge_properties = {k: v for k, v in record["r"].items()}
|
||||
result.edges.append(
|
||||
KnowledgeGraphEdge(
|
||||
id=edge_id,
|
||||
type="DIRECTED",
|
||||
source=source,
|
||||
target=target,
|
||||
properties=edge_properties,
|
||||
)
|
||||
)
|
||||
seen_edges.add(edge_id)
|
||||
else:
|
||||
# For specific label, use partial matching
|
||||
entity_name_label = node_label.strip('"')
|
||||
encoded_label = AGEStorage._encode_graph_label(entity_name_label)
|
||||
|
||||
# Find matching start nodes
|
||||
start_query = """
|
||||
MATCH (n:`{label}`)
|
||||
RETURN n
|
||||
"""
|
||||
start_params = {"label": encoded_label}
|
||||
start_nodes = await self._query(start_query, **start_params)
|
||||
|
||||
if not start_nodes:
|
||||
logger.warning(f"No nodes found with label '{entity_name_label}'!")
|
||||
return result
|
||||
|
||||
# Traverse graph from each start node
|
||||
for start_node_record in start_nodes:
|
||||
if "n" in start_node_record:
|
||||
# Use BFS to traverse graph
|
||||
query = """
|
||||
MATCH (start:`{label}`)
|
||||
CALL {
|
||||
MATCH path = (start)-[*0..{max_depth}]->(n)
|
||||
RETURN nodes(path) AS path_nodes, relationships(path) AS path_rels
|
||||
}
|
||||
RETURN DISTINCT path_nodes, path_rels
|
||||
"""
|
||||
params = {"label": encoded_label, "max_depth": max_depth}
|
||||
results = await self._query(query, **params)
|
||||
|
||||
# Extract nodes and edges from results
|
||||
for record in results:
|
||||
if "path_nodes" in record:
|
||||
# Process nodes
|
||||
for node in record["path_nodes"]:
|
||||
node_id = str(node.get("id", ""))
|
||||
if (
|
||||
node_id not in seen_nodes
|
||||
and len(seen_nodes) < max_graph_nodes
|
||||
):
|
||||
node_properties = {k: v for k, v in node.items()}
|
||||
node_label = node.get("label", "")
|
||||
result.nodes.append(
|
||||
KnowledgeGraphNode(
|
||||
id=node_id,
|
||||
labels=[node_label],
|
||||
properties=node_properties,
|
||||
)
|
||||
)
|
||||
seen_nodes.add(node_id)
|
||||
|
||||
if "path_rels" in record:
|
||||
# Process edges
|
||||
for rel in record["path_rels"]:
|
||||
source = str(rel.get("start_id", ""))
|
||||
target = str(rel.get("end_id", ""))
|
||||
edge_id = f"{source}-{target}"
|
||||
if edge_id not in seen_edges:
|
||||
edge_properties = {k: v for k, v in rel.items()}
|
||||
result.edges.append(
|
||||
KnowledgeGraphEdge(
|
||||
id=edge_id,
|
||||
type=rel.get("label", "DIRECTED"),
|
||||
source=source,
|
||||
target=target,
|
||||
properties=edge_properties,
|
||||
)
|
||||
)
|
||||
seen_edges.add(edge_id)
|
||||
|
||||
logger.info(
|
||||
f"Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}"
|
||||
)
|
||||
return result
|
||||
|
||||
async def index_done_callback(self) -> None:
|
||||
# AGES handles persistence automatically
|
||||
pass
|
||||
|
||||
async def drop(self) -> dict[str, str]:
|
||||
"""Drop the storage by removing all nodes and relationships in the graph.
|
||||
|
||||
Returns:
|
||||
dict[str, str]: Status of the operation with keys 'status' and 'message'
|
||||
"""
|
||||
try:
|
||||
query = """
|
||||
MATCH (n)
|
||||
DETACH DELETE n
|
||||
"""
|
||||
await self._query(query)
|
||||
logger.info(f"Successfully dropped all data from graph {self.graph_name}")
|
||||
return {"status": "success", "message": "graph data dropped"}
|
||||
except Exception as e:
|
||||
logger.error(f"Error dropping graph {self.graph_name}: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
|
@ -1,686 +0,0 @@
|
|||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import pipmaster as pm
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, final
|
||||
|
||||
from tenacity import (
|
||||
retry,
|
||||
retry_if_exception_type,
|
||||
stop_after_attempt,
|
||||
wait_exponential,
|
||||
)
|
||||
|
||||
from lightrag.types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
|
||||
from lightrag.utils import logger
|
||||
|
||||
from ..base import BaseGraphStorage
|
||||
|
||||
if not pm.is_installed("gremlinpython"):
|
||||
pm.install("gremlinpython")
|
||||
|
||||
from gremlin_python.driver import client, serializer # type: ignore
|
||||
from gremlin_python.driver.aiohttp.transport import AiohttpTransport # type: ignore
|
||||
from gremlin_python.driver.protocol import GremlinServerError # type: ignore
|
||||
|
||||
|
||||
@final
|
||||
@dataclass
|
||||
class GremlinStorage(BaseGraphStorage):
|
||||
@staticmethod
|
||||
def load_nx_graph(file_name):
|
||||
print("no preloading of graph with Gremlin in production")
|
||||
|
||||
def __init__(self, namespace, global_config, embedding_func):
|
||||
super().__init__(
|
||||
namespace=namespace,
|
||||
global_config=global_config,
|
||||
embedding_func=embedding_func,
|
||||
)
|
||||
|
||||
self._driver = None
|
||||
self._driver_lock = asyncio.Lock()
|
||||
|
||||
USER = os.environ.get("GREMLIN_USER", "")
|
||||
PASSWORD = os.environ.get("GREMLIN_PASSWORD", "")
|
||||
HOST = os.environ["GREMLIN_HOST"]
|
||||
PORT = int(os.environ["GREMLIN_PORT"])
|
||||
|
||||
# TraversalSource, a custom one has to be created manually,
|
||||
# default it "g"
|
||||
SOURCE = os.environ.get("GREMLIN_TRAVERSE_SOURCE", "g")
|
||||
|
||||
# All vertices will have graph={GRAPH} property, so that we can
|
||||
# have several logical graphs for one source
|
||||
GRAPH = GremlinStorage._to_value_map(
|
||||
os.environ.get("GREMLIN_GRAPH", "LightRAG")
|
||||
)
|
||||
|
||||
self.graph_name = GRAPH
|
||||
|
||||
self._driver = client.Client(
|
||||
f"ws://{HOST}:{PORT}/gremlin",
|
||||
SOURCE,
|
||||
username=USER,
|
||||
password=PASSWORD,
|
||||
message_serializer=serializer.GraphSONSerializersV3d0(),
|
||||
transport_factory=lambda: AiohttpTransport(call_from_event_loop=True),
|
||||
)
|
||||
|
||||
async def close(self):
|
||||
if self._driver:
|
||||
self._driver.close()
|
||||
self._driver = None
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
if self._driver:
|
||||
self._driver.close()
|
||||
|
||||
async def index_done_callback(self) -> None:
|
||||
# Gremlin handles persistence automatically
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _to_value_map(value: Any) -> str:
|
||||
"""Dump supported Python object as Gremlin valueMap"""
|
||||
json_str = json.dumps(value, ensure_ascii=False, sort_keys=False)
|
||||
parsed_str = json_str.replace("'", r"\'")
|
||||
|
||||
# walk over the string and replace curly brackets with square brackets
|
||||
# outside of strings, as well as replace double quotes with single quotes
|
||||
# and "deescape" double quotes inside of strings
|
||||
outside_str = True
|
||||
escaped = False
|
||||
remove_indices = []
|
||||
for i, c in enumerate(parsed_str):
|
||||
if escaped:
|
||||
# previous character was an "odd" backslash
|
||||
escaped = False
|
||||
if c == '"':
|
||||
# we want to "deescape" double quotes: store indices to delete
|
||||
remove_indices.insert(0, i - 1)
|
||||
elif c == "\\":
|
||||
escaped = True
|
||||
elif c == '"':
|
||||
outside_str = not outside_str
|
||||
parsed_str = parsed_str[:i] + "'" + parsed_str[i + 1 :]
|
||||
elif c == "{" and outside_str:
|
||||
parsed_str = parsed_str[:i] + "[" + parsed_str[i + 1 :]
|
||||
elif c == "}" and outside_str:
|
||||
parsed_str = parsed_str[:i] + "]" + parsed_str[i + 1 :]
|
||||
for idx in remove_indices:
|
||||
parsed_str = parsed_str[:idx] + parsed_str[idx + 1 :]
|
||||
return parsed_str
|
||||
|
||||
@staticmethod
|
||||
def _convert_properties(properties: Dict[str, Any]) -> str:
|
||||
"""Create chained .property() commands from properties dict"""
|
||||
props = []
|
||||
for k, v in properties.items():
|
||||
prop_name = GremlinStorage._to_value_map(k)
|
||||
props.append(f".property({prop_name}, {GremlinStorage._to_value_map(v)})")
|
||||
return "".join(props)
|
||||
|
||||
@staticmethod
|
||||
def _fix_name(name: str) -> str:
|
||||
"""Strip double quotes and format as a proper field name"""
|
||||
name = GremlinStorage._to_value_map(name.strip('"').replace(r"\'", "'"))
|
||||
|
||||
return name
|
||||
|
||||
async def _query(self, query: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Query the Gremlin graph
|
||||
|
||||
Args:
|
||||
query (str): a query to be executed
|
||||
|
||||
Returns:
|
||||
List[Dict[str, Any]]: a list of dictionaries containing the result set
|
||||
"""
|
||||
|
||||
result = list(await asyncio.wrap_future(self._driver.submit_async(query)))
|
||||
if result:
|
||||
result = result[0]
|
||||
|
||||
return result
|
||||
|
||||
async def has_node(self, node_id: str) -> bool:
|
||||
entity_name = GremlinStorage._fix_name(node_id)
|
||||
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.has('entity_name', {entity_name})
|
||||
.limit(1)
|
||||
.count()
|
||||
.project('has_node')
|
||||
.by(__.choose(__.is(gt(0)), constant(true), constant(false)))
|
||||
"""
|
||||
result = await self._query(query)
|
||||
logger.debug(
|
||||
"{%s}:query:{%s}:result:{%s}",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
query,
|
||||
result[0]["has_node"],
|
||||
)
|
||||
|
||||
return result[0]["has_node"]
|
||||
|
||||
async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
|
||||
entity_name_source = GremlinStorage._fix_name(source_node_id)
|
||||
entity_name_target = GremlinStorage._fix_name(target_node_id)
|
||||
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.has('entity_name', {entity_name_source})
|
||||
.outE()
|
||||
.inV().has('graph', {self.graph_name})
|
||||
.has('entity_name', {entity_name_target})
|
||||
.limit(1)
|
||||
.count()
|
||||
.project('has_edge')
|
||||
.by(__.choose(__.is(gt(0)), constant(true), constant(false)))
|
||||
"""
|
||||
result = await self._query(query)
|
||||
logger.debug(
|
||||
"{%s}:query:{%s}:result:{%s}",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
query,
|
||||
result[0]["has_edge"],
|
||||
)
|
||||
|
||||
return result[0]["has_edge"]
|
||||
|
||||
async def get_node(self, node_id: str) -> dict[str, str] | None:
|
||||
entity_name = GremlinStorage._fix_name(node_id)
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.has('entity_name', {entity_name})
|
||||
.limit(1)
|
||||
.project('properties')
|
||||
.by(elementMap())
|
||||
"""
|
||||
result = await self._query(query)
|
||||
if result:
|
||||
node = result[0]
|
||||
node_dict = node["properties"]
|
||||
logger.debug(
|
||||
"{%s}: query: {%s}, result: {%s}",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
query.format,
|
||||
node_dict,
|
||||
)
|
||||
return node_dict
|
||||
|
||||
async def node_degree(self, node_id: str) -> int:
|
||||
entity_name = GremlinStorage._fix_name(node_id)
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.has('entity_name', {entity_name})
|
||||
.outE()
|
||||
.inV().has('graph', {self.graph_name})
|
||||
.count()
|
||||
.project('total_edge_count')
|
||||
.by()
|
||||
"""
|
||||
result = await self._query(query)
|
||||
edge_count = result[0]["total_edge_count"]
|
||||
|
||||
logger.debug(
|
||||
"{%s}:query:{%s}:result:{%s}",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
query,
|
||||
edge_count,
|
||||
)
|
||||
|
||||
return edge_count
|
||||
|
||||
async def edge_degree(self, src_id: str, tgt_id: str) -> int:
|
||||
src_degree = await self.node_degree(src_id)
|
||||
trg_degree = await self.node_degree(tgt_id)
|
||||
|
||||
# Convert None to 0 for addition
|
||||
src_degree = 0 if src_degree is None else src_degree
|
||||
trg_degree = 0 if trg_degree is None else trg_degree
|
||||
|
||||
degrees = int(src_degree) + int(trg_degree)
|
||||
logger.debug(
|
||||
"{%s}:query:src_Degree+trg_degree:result:{%s}",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
degrees,
|
||||
)
|
||||
return degrees
|
||||
|
||||
async def get_edge(
|
||||
self, source_node_id: str, target_node_id: str
|
||||
) -> dict[str, str] | None:
|
||||
entity_name_source = GremlinStorage._fix_name(source_node_id)
|
||||
entity_name_target = GremlinStorage._fix_name(target_node_id)
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.has('entity_name', {entity_name_source})
|
||||
.outE()
|
||||
.inV().has('graph', {self.graph_name})
|
||||
.has('entity_name', {entity_name_target})
|
||||
.limit(1)
|
||||
.project('edge_properties')
|
||||
.by(__.bothE().elementMap())
|
||||
"""
|
||||
result = await self._query(query)
|
||||
if result:
|
||||
edge_properties = result[0]["edge_properties"]
|
||||
logger.debug(
|
||||
"{%s}:query:{%s}:result:{%s}",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
query,
|
||||
edge_properties,
|
||||
)
|
||||
return edge_properties
|
||||
|
||||
async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None:
|
||||
node_name = GremlinStorage._fix_name(source_node_id)
|
||||
query = f"""g
|
||||
.E()
|
||||
.filter(
|
||||
__.or(
|
||||
__.outV().has('graph', {self.graph_name})
|
||||
.has('entity_name', {node_name}),
|
||||
__.inV().has('graph', {self.graph_name})
|
||||
.has('entity_name', {node_name})
|
||||
)
|
||||
)
|
||||
.project('source_name', 'target_name')
|
||||
.by(__.outV().values('entity_name'))
|
||||
.by(__.inV().values('entity_name'))
|
||||
"""
|
||||
result = await self._query(query)
|
||||
edges = [(res["source_name"], res["target_name"]) for res in result]
|
||||
|
||||
return edges
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(10),
|
||||
wait=wait_exponential(multiplier=1, min=4, max=10),
|
||||
retry=retry_if_exception_type((GremlinServerError,)),
|
||||
)
|
||||
async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
|
||||
"""
|
||||
Upsert a node in the Gremlin graph.
|
||||
|
||||
Args:
|
||||
node_id: The unique identifier for the node (used as name)
|
||||
node_data: Dictionary of node properties
|
||||
"""
|
||||
name = GremlinStorage._fix_name(node_id)
|
||||
properties = GremlinStorage._convert_properties(node_data)
|
||||
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.has('entity_name', {name})
|
||||
.fold()
|
||||
.coalesce(
|
||||
__.unfold(),
|
||||
__.addV('ENTITY')
|
||||
.property('graph', {self.graph_name})
|
||||
.property('entity_name', {name})
|
||||
)
|
||||
{properties}
|
||||
"""
|
||||
|
||||
try:
|
||||
await self._query(query)
|
||||
logger.debug(
|
||||
"Upserted node with name {%s} and properties: {%s}",
|
||||
name,
|
||||
properties,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Error during upsert: {%s}", e)
|
||||
raise
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(10),
|
||||
wait=wait_exponential(multiplier=1, min=4, max=10),
|
||||
retry=retry_if_exception_type((GremlinServerError,)),
|
||||
)
|
||||
async def upsert_edge(
|
||||
self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
|
||||
) -> None:
|
||||
"""
|
||||
Upsert an edge and its properties between two nodes identified by their names.
|
||||
|
||||
Args:
|
||||
source_node_id (str): Name of the source node (used as identifier)
|
||||
target_node_id (str): Name of the target node (used as identifier)
|
||||
edge_data (dict): Dictionary of properties to set on the edge
|
||||
"""
|
||||
source_node_name = GremlinStorage._fix_name(source_node_id)
|
||||
target_node_name = GremlinStorage._fix_name(target_node_id)
|
||||
edge_properties = GremlinStorage._convert_properties(edge_data)
|
||||
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.has('entity_name', {source_node_name}).as('source')
|
||||
.V().has('graph', {self.graph_name})
|
||||
.has('entity_name', {target_node_name}).as('target')
|
||||
.coalesce(
|
||||
__.select('source').outE('DIRECTED').where(__.inV().as('target')),
|
||||
__.select('source').addE('DIRECTED').to(__.select('target'))
|
||||
)
|
||||
.property('graph', {self.graph_name})
|
||||
{edge_properties}
|
||||
"""
|
||||
try:
|
||||
await self._query(query)
|
||||
logger.debug(
|
||||
"Upserted edge from {%s} to {%s} with properties: {%s}",
|
||||
source_node_name,
|
||||
target_node_name,
|
||||
edge_properties,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Error during edge upsert: {%s}", e)
|
||||
raise
|
||||
|
||||
async def delete_node(self, node_id: str) -> None:
|
||||
"""Delete a node with the specified entity_name
|
||||
|
||||
Args:
|
||||
node_id: The entity_name of the node to delete
|
||||
"""
|
||||
entity_name = GremlinStorage._fix_name(node_id)
|
||||
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.has('entity_name', {entity_name})
|
||||
.drop()
|
||||
"""
|
||||
try:
|
||||
await self._query(query)
|
||||
logger.debug(
|
||||
"{%s}: Deleted node with entity_name '%s'",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
entity_name,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during node deletion: {str(e)}")
|
||||
raise
|
||||
|
||||
async def get_all_labels(self) -> list[str]:
|
||||
"""
|
||||
Get all node entity_names in the graph
|
||||
Returns:
|
||||
[entity_name1, entity_name2, ...] # Alphabetically sorted entity_name list
|
||||
"""
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.values('entity_name')
|
||||
.dedup()
|
||||
.order()
|
||||
"""
|
||||
try:
|
||||
result = await self._query(query)
|
||||
labels = result if result else []
|
||||
logger.debug(
|
||||
"{%s}: Retrieved %d labels",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
len(labels),
|
||||
)
|
||||
return labels
|
||||
except Exception as e:
|
||||
logger.error(f"Error retrieving labels: {str(e)}")
|
||||
return []
|
||||
|
||||
async def get_knowledge_graph(
|
||||
self, node_label: str, max_depth: int = 5
|
||||
) -> KnowledgeGraph:
|
||||
"""
|
||||
Retrieve a connected subgraph of nodes where the entity_name includes the specified `node_label`.
|
||||
Maximum number of nodes is constrained by the environment variable `MAX_GRAPH_NODES` (default: 1000).
|
||||
|
||||
Args:
|
||||
node_label: Entity name of the starting node
|
||||
max_depth: Maximum depth of the subgraph
|
||||
|
||||
Returns:
|
||||
KnowledgeGraph object containing nodes and edges
|
||||
"""
|
||||
result = KnowledgeGraph()
|
||||
seen_nodes = set()
|
||||
seen_edges = set()
|
||||
|
||||
# Get maximum number of graph nodes from environment variable, default is 1000
|
||||
MAX_GRAPH_NODES = int(os.getenv("MAX_GRAPH_NODES", 1000))
|
||||
|
||||
entity_name = GremlinStorage._fix_name(node_label)
|
||||
|
||||
# Handle special case for "*" label
|
||||
if node_label == "*":
|
||||
# For "*", get all nodes and their edges (limited by MAX_GRAPH_NODES)
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.limit({MAX_GRAPH_NODES})
|
||||
.elementMap()
|
||||
"""
|
||||
nodes_result = await self._query(query)
|
||||
|
||||
# Add nodes to result
|
||||
for node_data in nodes_result:
|
||||
node_id = node_data.get("entity_name", str(node_data.get("id", "")))
|
||||
if str(node_id) in seen_nodes:
|
||||
continue
|
||||
|
||||
# Create node with properties
|
||||
node_properties = {
|
||||
k: v for k, v in node_data.items() if k not in ["id", "label"]
|
||||
}
|
||||
|
||||
result.nodes.append(
|
||||
KnowledgeGraphNode(
|
||||
id=str(node_id),
|
||||
labels=[str(node_id)],
|
||||
properties=node_properties,
|
||||
)
|
||||
)
|
||||
seen_nodes.add(str(node_id))
|
||||
|
||||
# Get and add edges
|
||||
if nodes_result:
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.limit({MAX_GRAPH_NODES})
|
||||
.outE()
|
||||
.inV().has('graph', {self.graph_name})
|
||||
.limit({MAX_GRAPH_NODES})
|
||||
.path()
|
||||
.by(elementMap())
|
||||
.by(elementMap())
|
||||
.by(elementMap())
|
||||
"""
|
||||
edges_result = await self._query(query)
|
||||
|
||||
for path in edges_result:
|
||||
if len(path) >= 3: # source -> edge -> target
|
||||
source = path[0]
|
||||
edge_data = path[1]
|
||||
target = path[2]
|
||||
|
||||
source_id = source.get("entity_name", str(source.get("id", "")))
|
||||
target_id = target.get("entity_name", str(target.get("id", "")))
|
||||
|
||||
edge_id = f"{source_id}-{target_id}"
|
||||
if edge_id in seen_edges:
|
||||
continue
|
||||
|
||||
# Create edge with properties
|
||||
edge_properties = {
|
||||
k: v
|
||||
for k, v in edge_data.items()
|
||||
if k not in ["id", "label"]
|
||||
}
|
||||
|
||||
result.edges.append(
|
||||
KnowledgeGraphEdge(
|
||||
id=edge_id,
|
||||
type="DIRECTED",
|
||||
source=str(source_id),
|
||||
target=str(target_id),
|
||||
properties=edge_properties,
|
||||
)
|
||||
)
|
||||
seen_edges.add(edge_id)
|
||||
else:
|
||||
# Search for specific node and get its neighborhood
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.has('entity_name', {entity_name})
|
||||
.repeat(__.both().simplePath().dedup())
|
||||
.times({max_depth})
|
||||
.emit()
|
||||
.dedup()
|
||||
.limit({MAX_GRAPH_NODES})
|
||||
.elementMap()
|
||||
"""
|
||||
nodes_result = await self._query(query)
|
||||
|
||||
# Add nodes to result
|
||||
for node_data in nodes_result:
|
||||
node_id = node_data.get("entity_name", str(node_data.get("id", "")))
|
||||
if str(node_id) in seen_nodes:
|
||||
continue
|
||||
|
||||
# Create node with properties
|
||||
node_properties = {
|
||||
k: v for k, v in node_data.items() if k not in ["id", "label"]
|
||||
}
|
||||
|
||||
result.nodes.append(
|
||||
KnowledgeGraphNode(
|
||||
id=str(node_id),
|
||||
labels=[str(node_id)],
|
||||
properties=node_properties,
|
||||
)
|
||||
)
|
||||
seen_nodes.add(str(node_id))
|
||||
|
||||
# Get edges between the nodes in the result
|
||||
if nodes_result:
|
||||
node_ids = [
|
||||
n.get("entity_name", str(n.get("id", ""))) for n in nodes_result
|
||||
]
|
||||
node_ids_query = ", ".join(
|
||||
[GremlinStorage._to_value_map(nid) for nid in node_ids]
|
||||
)
|
||||
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.has('entity_name', within({node_ids_query}))
|
||||
.outE()
|
||||
.where(inV().has('graph', {self.graph_name})
|
||||
.has('entity_name', within({node_ids_query})))
|
||||
.path()
|
||||
.by(elementMap())
|
||||
.by(elementMap())
|
||||
.by(elementMap())
|
||||
"""
|
||||
edges_result = await self._query(query)
|
||||
|
||||
for path in edges_result:
|
||||
if len(path) >= 3: # source -> edge -> target
|
||||
source = path[0]
|
||||
edge_data = path[1]
|
||||
target = path[2]
|
||||
|
||||
source_id = source.get("entity_name", str(source.get("id", "")))
|
||||
target_id = target.get("entity_name", str(target.get("id", "")))
|
||||
|
||||
edge_id = f"{source_id}-{target_id}"
|
||||
if edge_id in seen_edges:
|
||||
continue
|
||||
|
||||
# Create edge with properties
|
||||
edge_properties = {
|
||||
k: v
|
||||
for k, v in edge_data.items()
|
||||
if k not in ["id", "label"]
|
||||
}
|
||||
|
||||
result.edges.append(
|
||||
KnowledgeGraphEdge(
|
||||
id=edge_id,
|
||||
type="DIRECTED",
|
||||
source=str(source_id),
|
||||
target=str(target_id),
|
||||
properties=edge_properties,
|
||||
)
|
||||
)
|
||||
seen_edges.add(edge_id)
|
||||
|
||||
logger.info(
|
||||
"Subgraph query successful | Node count: %d | Edge count: %d",
|
||||
len(result.nodes),
|
||||
len(result.edges),
|
||||
)
|
||||
return result
|
||||
|
||||
async def remove_nodes(self, nodes: list[str]):
|
||||
"""Delete multiple nodes
|
||||
|
||||
Args:
|
||||
nodes: List of node entity_names to be deleted
|
||||
"""
|
||||
for node in nodes:
|
||||
await self.delete_node(node)
|
||||
|
||||
async def remove_edges(self, edges: list[tuple[str, str]]):
|
||||
"""Delete multiple edges
|
||||
|
||||
Args:
|
||||
edges: List of edges to be deleted, each edge is a (source, target) tuple
|
||||
"""
|
||||
for source, target in edges:
|
||||
entity_name_source = GremlinStorage._fix_name(source)
|
||||
entity_name_target = GremlinStorage._fix_name(target)
|
||||
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.has('entity_name', {entity_name_source})
|
||||
.outE()
|
||||
.where(inV().has('graph', {self.graph_name})
|
||||
.has('entity_name', {entity_name_target}))
|
||||
.drop()
|
||||
"""
|
||||
try:
|
||||
await self._query(query)
|
||||
logger.debug(
|
||||
"{%s}: Deleted edge from '%s' to '%s'",
|
||||
inspect.currentframe().f_code.co_name,
|
||||
entity_name_source,
|
||||
entity_name_target,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during edge deletion: {str(e)}")
|
||||
raise
|
||||
|
||||
async def drop(self) -> dict[str, str]:
|
||||
"""Drop the storage by removing all nodes and relationships in the graph.
|
||||
|
||||
This function deletes all nodes with the specified graph name property,
|
||||
which automatically removes all associated edges.
|
||||
|
||||
Returns:
|
||||
dict[str, str]: Status of the operation with keys 'status' and 'message'
|
||||
"""
|
||||
try:
|
||||
query = f"""g
|
||||
.V().has('graph', {self.graph_name})
|
||||
.drop()
|
||||
"""
|
||||
await self._query(query)
|
||||
logger.info(f"Successfully dropped all data from graph {self.graph_name}")
|
||||
return {"status": "success", "message": "graph data dropped"}
|
||||
except Exception as e:
|
||||
logger.error(f"Error dropping graph {self.graph_name}: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
File diff suppressed because it is too large
Load diff
Loading…
Add table
Reference in a new issue