lightrag/kg/gremlin_impl.py

import asyncio
import inspect
import json
import os
from dataclasses import dataclass
from typing import Any, Dict, List, Tuple, Union

from gremlin_python.driver import client, serializer
from gremlin_python.driver.aiohttp.transport import AiohttpTransport
from gremlin_python.driver.protocol import GremlinServerError
from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from lightrag.utils import logger

from ..base import BaseGraphStorage


@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["GREMLIN_GRAPH"])

        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),
        )

    def __post_init__(self):
        self._node_embed_algorithms = {
            "node2vec": self._node2vec_embed,
        }

    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):
        print("KG successfully indexed.")

    @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) -> Union[dict, 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
    ) -> Union[dict, None]:
        """
        Find all edges between nodes of two given names

        Args:
            source_node_id (str): Name of the source nodes
            target_node_id (str): Name of the target nodes

        Returns:
            dict|None: Dict of found edge properties, or None if not found
        """
        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]]:
        """
        Retrieves all edges (relationships) for a particular node identified by its name.
        :return: List of tuples containing edge sources and targets
        """
        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, Any]):
        """
        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, Any]
    ):
        """
        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 _node2vec_embed(self):
        print("Implemented but never called.")