cognee 0.2.1.dev7__py3-none-any.whl → 0.2.2.dev1__py3-none-any.whl

cognee/infrastructure/databases/graph/neptune_driver/adapter.py

@@ -0,0 +1,1427 @@
+"""Neptune Analytics Adapter for Graph Database"""
+
+import json
+from typing import Optional, Any, List, Dict, Type, Tuple
+from uuid import UUID
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.databases.graph.graph_db_interface import (
+    GraphDBInterface,
+    record_graph_changes,
+    NodeData,
+    EdgeData,
+    Node,
+)
+from cognee.modules.storage.utils import JSONEncoder
+from cognee.infrastructure.engine import DataPoint
+from botocore.config import Config
+
+from .exceptions import (
+    NeptuneAnalyticsConfigurationError,
+)
+from .neptune_utils import (
+    validate_graph_id,
+    validate_aws_region,
+    build_neptune_config,
+    format_neptune_error,
+)
+
+logger = get_logger("NeptuneGraphDB")
+
+try:
+    from langchain_aws import NeptuneAnalyticsGraph
+
+    LANGCHAIN_AWS_AVAILABLE = True
+except ImportError:
+    logger.warning("langchain_aws not available. Neptune Analytics functionality will be limited.")
+    LANGCHAIN_AWS_AVAILABLE = False
+
+NEPTUNE_ENDPOINT_URL = "neptune-graph://"
+
+
+class NeptuneGraphDB(GraphDBInterface):
+    """
+    Adapter for interacting with Amazon Neptune Analytics graph store.
+    This class provides methods for querying, adding, deleting nodes and edges using the aws_langchain library.
+    """
+
+    _GRAPH_NODE_LABEL = "COGNEE_NODE"
+
+    def __init__(
+        self,
+        graph_id: str,
+        region: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+        aws_session_token: Optional[str] = None,
+    ):
+        """
+        Initialize the Neptune Analytics adapter.
+
+        Parameters:
+        -----------
+            - graph_id (str): The Neptune Analytics graph identifier
+            - region (Optional[str]): AWS region where the graph is located (default: us-east-1)
+            - aws_access_key_id (Optional[str]): AWS access key ID
+            - aws_secret_access_key (Optional[str]): AWS secret access key
+            - aws_session_token (Optional[str]): AWS session token for temporary credentials
+
+        Raises:
+        -------
+            - NeptuneAnalyticsConfigurationError: If configuration parameters are invalid
+        """
+        # validate import
+        if not LANGCHAIN_AWS_AVAILABLE:
+            raise ImportError(
+                "langchain_aws is not available. Please install it to use Neptune Analytics."
+            )
+
+        # Validate configuration
+        if not validate_graph_id(graph_id):
+            raise NeptuneAnalyticsConfigurationError(message=f'Invalid graph ID: "{graph_id}"')
+
+        if region and not validate_aws_region(region):
+            raise NeptuneAnalyticsConfigurationError(message=f'Invalid AWS region: "{region}"')
+
+        self.graph_id = graph_id
+        self.region = region
+        self.aws_access_key_id = aws_access_key_id
+        self.aws_secret_access_key = aws_secret_access_key
+        self.aws_session_token = aws_session_token
+
+        # Build configuration
+        self.config = build_neptune_config(
+            graph_id=self.graph_id,
+            region=self.region,
+            aws_access_key_id=self.aws_access_key_id,
+            aws_secret_access_key=self.aws_secret_access_key,
+            aws_session_token=self.aws_session_token,
+        )
+
+        # Initialize Neptune Analytics client using langchain_aws
+        self._client: NeptuneAnalyticsGraph = self._initialize_client()
+        logger.info(
+            f'Initialized Neptune Analytics adapter for graph: "{graph_id}" in region: "{self.region}"'
+        )
+
+    def _initialize_client(self) -> Optional[NeptuneAnalyticsGraph]:
+        """
+        Initialize the Neptune Analytics client using langchain_aws.
+
+        Returns:
+        --------
+            - Optional[Any]: The Neptune Analytics client or None if not available
+        """
+        try:
+            # Initialize the Neptune Analytics Graph client
+            client_config = {
+                "graph_identifier": self.graph_id,
+                "config": Config(user_agent_appid="Cognee"),
+            }
+            # Add AWS credentials if provided
+            if self.region:
+                client_config["region_name"] = self.region
+            if self.aws_access_key_id:
+                client_config["aws_access_key_id"] = self.aws_access_key_id
+            if self.aws_secret_access_key:
+                client_config["aws_secret_access_key"] = self.aws_secret_access_key
+            if self.aws_session_token:
+                client_config["aws_session_token"] = self.aws_session_token
+
+            client = NeptuneAnalyticsGraph(**client_config)
+            logger.info("Successfully initialized Neptune Analytics client")
+            return client
+
+        except Exception as e:
+            raise NeptuneAnalyticsConfigurationError(
+                message=f"Failed to initialize Neptune Analytics client: {format_neptune_error(e)}"
+            ) from e
+
+    @staticmethod
+    def _serialize_properties(properties: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Serialize properties for Neptune Analytics storage.
+        Parameters:
+        -----------
+            - properties (Dict[str, Any]): Properties to serialize.
+        Returns:
+        --------
+            - Dict[str, Any]: Serialized properties.
+        """
+        serialized_properties = {}
+
+        for property_key, property_value in properties.items():
+            if isinstance(property_value, UUID):
+                serialized_properties[property_key] = str(property_value)
+                continue
+
+            if isinstance(property_value, dict) or isinstance(property_value, list):
+                serialized_properties[property_key] = json.dumps(property_value, cls=JSONEncoder)
+                continue
+
+            serialized_properties[property_key] = property_value
+
+        return serialized_properties
+
+    async def query(self, query: str, params: Optional[Dict[str, Any]] = None) -> List[Any]:
+        """
+        Execute a query against the Neptune Analytics database and return the results.
+
+        Parameters:
+        -----------
+            - query (str): The query string to execute against the database.
+            - params (Optional[Dict[str, Any]]): A dictionary of parameters to be used in the query.
+
+        Returns:
+        --------
+            - List[Any]: A list of results from the query execution.
+        """
+        try:
+            # Execute the query using the Neptune Analytics client
+            # The langchain_aws NeptuneAnalyticsGraph supports openCypher queries
+            if params is None:
+                params = {}
+            logger.debug(f"executing na query:\nquery={query}\n")
+            result = self._client.query(query, params)
+
+            # Convert the result to list format expected by the interface
+            if isinstance(result, list):
+                return result
+            elif isinstance(result, dict):
+                return [result]
+            else:
+                return [{"result": result}]
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Neptune Analytics query failed: {error_msg}")
+            raise Exception(f"Query execution failed: {error_msg}") from e
+
+    async def add_node(self, node: DataPoint) -> None:
+        """
+        Add a single node with specified properties to the graph.
+
+        Parameters:
+        -----------
+            - node (DataPoint): The DataPoint object to be added to the graph.
+        """
+        try:
+            # Prepare node properties with the ID and graph type
+            serialized_properties = self._serialize_properties(node.model_dump())
+
+            query = f"""
+            MERGE (n:{self._GRAPH_NODE_LABEL} {{`~id`: $node_id}})
+            ON CREATE SET n = $properties, n.updated_at = timestamp()
+            ON MATCH SET n += $properties, n.updated_at = timestamp()
+            RETURN n
+            """
+
+            params = {
+                "node_id": str(node.id),
+                "properties": serialized_properties,
+            }
+
+            result = await self.query(query, params)
+            logger.debug(f"Successfully added/updated node: {node.id}")
+            logger.debug(f"Successfully added/updated node: {str(result)}")
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to add node {node.id}: {error_msg}")
+            raise Exception(f"Failed to add node: {error_msg}") from e
+
+    @record_graph_changes
+    async def add_nodes(self, nodes: List[DataPoint]) -> None:
+        """
+        Add multiple nodes to the graph in a single operation.
+
+        Parameters:
+        -----------
+            - nodes (List[DataPoint]): A list of DataPoint objects to be added to the graph.
+        """
+        if not nodes:
+            logger.debug("No nodes to add")
+            return
+
+        try:
+            # Build bulk node creation query using UNWIND
+            query = f"""
+            UNWIND $nodes AS node
+            MERGE (n:{self._GRAPH_NODE_LABEL} {{`~id`: node.node_id}})
+            ON CREATE SET n = node.properties, n.updated_at = timestamp()
+            ON MATCH SET n += node.properties, n.updated_at = timestamp()
+            RETURN count(n) AS nodes_processed
+            """
+
+            # Prepare node data for bulk operation
+            params = {
+                "nodes": [
+                    {
+                        "node_id": str(node.id),
+                        "properties": self._serialize_properties(node.model_dump()),
+                    }
+                    for node in nodes
+                ]
+            }
+            result = await self.query(query, params)
+
+            processed_count = result[0].get("nodes_processed", 0) if result else 0
+            logger.debug(f"Successfully processed {processed_count} nodes in bulk operation")
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to add nodes in bulk: {error_msg}")
+            # Fallback to individual node creation
+            logger.info("Falling back to individual node creation")
+            for node in nodes:
+                try:
+                    await self.add_node(node)
+                except Exception as node_error:
+                    logger.error(
+                        f"Failed to add individual node {node.id}: {format_neptune_error(node_error)}"
+                    )
+                    continue
+
+    async def delete_node(self, node_id: str) -> None:
+        """
+        Delete a specified node from the graph by its ID.
+
+        Parameters:
+        -----------
+            - node_id (str): Unique identifier for the node to delete.
+        """
+        try:
+            # Build openCypher query to delete the node and all its relationships
+            query = f"""
+            MATCH (n:{self._GRAPH_NODE_LABEL})
+            WHERE id(n) = $node_id
+            DETACH DELETE n
+            """
+
+            params = {"node_id": node_id}
+
+            await self.query(query, params)
+            logger.debug(f"Successfully deleted node: {node_id}")
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to delete node {node_id}: {error_msg}")
+            raise Exception(f"Failed to delete node: {error_msg}") from e
+
+    async def delete_nodes(self, node_ids: List[str]) -> None:
+        """
+        Delete multiple nodes from the graph by their identifiers.
+
+        Parameters:
+        -----------
+            - node_ids (List[str]): A list of unique identifiers for the nodes to delete.
+        """
+        if not node_ids:
+            logger.debug("No nodes to delete")
+            return
+
+        try:
+            # Build bulk node deletion query using UNWIND
+            query = f"""
+            UNWIND $node_ids AS node_id
+            MATCH (n:{self._GRAPH_NODE_LABEL})
+            WHERE id(n) = node_id
+            DETACH DELETE n
+            """
+
+            params = {"node_ids": node_ids}
+            await self.query(query, params)
+            logger.debug(f"Successfully deleted {len(node_ids)} nodes in bulk operation")
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to delete nodes in bulk: {error_msg}")
+            # Fallback to individual node deletion
+            logger.info("Falling back to individual node deletion")
+            for node_id in node_ids:
+                try:
+                    await self.delete_node(node_id)
+                except Exception as node_error:
+                    logger.error(
+                        f"Failed to delete individual node {node_id}: {format_neptune_error(node_error)}"
+                    )
+                    continue
+
+    async def get_node(self, node_id: str) -> Optional[NodeData]:
+        """
+        Retrieve a single node from the graph using its ID.
+
+        Parameters:
+        -----------
+            - node_id (str): Unique identifier of the node to retrieve.
+
+        Returns:
+        --------
+            - Optional[NodeData]: The node data if found, None otherwise.
+        """
+        try:
+            # Build openCypher query to retrieve the node
+            query = f"""
+            MATCH (n:{self._GRAPH_NODE_LABEL})
+            WHERE id(n) = $node_id
+            RETURN n
+            """
+            params = {"node_id": node_id}
+
+            result = await self.query(query, params)
+
+            if result and len(result) == 1:
+                # Extract node properties from the result
+                logger.debug(f"Successfully retrieved node: {node_id}")
+                return result[0]["n"]
+            else:
+                if not result:
+                    logger.debug(f"Node not found: {node_id}")
+                elif len(result) > 1:
+                    logger.debug(f"Only one node expected, multiple returned: {node_id}")
+                return None
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to get node {node_id}: {error_msg}")
+            raise Exception(f"Failed to get node: {error_msg}") from e
+
+    async def get_nodes(self, node_ids: List[str]) -> List[NodeData]:
+        """
+        Retrieve multiple nodes from the graph using their IDs.
+
+        Parameters:
+        -----------
+            - node_ids (List[str]): A list of unique identifiers for the nodes to retrieve.
+
+        Returns:
+        --------
+            - List[NodeData]: A list of node data for the found nodes.
+        """
+        if not node_ids:
+            logger.debug("No node IDs provided")
+            return []
+
+        try:
+            # Build bulk node-retrieval OpenCypher query using UNWIND
+            query = f"""
+            UNWIND $node_ids AS node_id
+            MATCH (n:{self._GRAPH_NODE_LABEL})
+            WHERE id(n) = node_id
+            RETURN n
+            """
+
+            params = {"node_ids": node_ids}
+            result = await self.query(query, params)
+
+            # Extract node data from results
+            nodes = [record["n"] for record in result]
+
+            logger.debug(
+                f"Successfully retrieved {len(nodes)} nodes out of {len(node_ids)} requested"
+            )
+            return nodes
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to get nodes in bulk: {error_msg}")
+            # Fallback to individual node retrieval
+            logger.info("Falling back to individual node retrieval")
+            nodes = []
+            for node_id in node_ids:
+                try:
+                    node_data = await self.get_node(node_id)
+                    if node_data:
+                        nodes.append(node_data)
+                except Exception as node_error:
+                    logger.error(
+                        f"Failed to get individual node {node_id}: {format_neptune_error(node_error)}"
+                    )
+                    continue
+            return nodes
+
+    async def extract_node(self, node_id: str):
+        """
+        Retrieve a single node based on its ID.
+
+        Parameters:
+        -----------
+
+            - node_id (str): The ID of the node to retrieve.
+
+        Returns:
+        --------
+
+            - Optional[Dict[str, Any]]: The requested node as a dictionary, or None if it does
+              not exist.
+        """
+        results = await self.extract_nodes([node_id])
+
+        return results[0] if len(results) > 0 else None
+
+    async def extract_nodes(self, node_ids: List[str]):
+        """
+        Retrieve multiple nodes from the database by their IDs.
+
+        Parameters:
+        -----------
+
+            - node_ids (List[str]): A list of IDs for the nodes to retrieve.
+
+        Returns:
+        --------
+
+            A list of nodes represented as dictionaries.
+        """
+        query = f"""
+        UNWIND $node_ids AS id
+        MATCH (node :{self._GRAPH_NODE_LABEL}) WHERE id(node) = id
+        RETURN node"""
+
+        params = {"node_ids": node_ids}
+
+        results = await self.query(query, params)
+
+        return [result["node"] for result in results]
+
+    async def add_edge(
+        self,
+        source_id: str,
+        target_id: str,
+        relationship_name: str,
+        properties: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """
+        Create a new edge between two nodes in the graph.
+
+        Parameters:
+        -----------
+            - source_id (str): The unique identifier of the source node.
+            - target_id (str): The unique identifier of the target node.
+            - relationship_name (str): The name of the relationship to be established by the edge.
+            - properties (Optional[Dict[str, Any]]): Optional dictionary of properties associated with the edge.
+        """
+        try:
+            # Build openCypher query to create the edge
+            # First ensure both nodes exist, then create the relationship
+
+            # Prepare edge properties
+            edge_props = properties or {}
+            serialized_properties = self._serialize_properties(edge_props)
+
+            query = f"""
+            MATCH (source:{self._GRAPH_NODE_LABEL})
+            WHERE id(source) = $source_id
+            MATCH (target:{self._GRAPH_NODE_LABEL})
+            WHERE id(target) = $target_id
+            MERGE (source)-[r:{relationship_name}]->(target)
+            ON CREATE SET r = $properties, r.updated_at = timestamp()
+            ON MATCH SET r = $properties, r.updated_at = timestamp()
+            RETURN r
+            """
+
+            params = {
+                "source_id": source_id,
+                "target_id": target_id,
+                "properties": serialized_properties,
+            }
+            await self.query(query, params)
+            logger.debug(
+                f"Successfully added edge: {source_id} -[{relationship_name}]-> {target_id}"
+            )
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to add edge {source_id} -> {target_id}: {error_msg}")
+            raise Exception(f"Failed to add edge: {error_msg}") from e
+
+    @record_graph_changes
+    async def add_edges(self, edges: List[Tuple[str, str, str, Optional[Dict[str, Any]]]]) -> None:
+        """
+        Add multiple edges to the graph in a single operation.
+
+        Parameters:
+        -----------
+            - edges (List[EdgeData]): A list of EdgeData objects representing edges to be added.
+        """
+        if not edges:
+            logger.debug("No edges to add")
+            return
+
+        edges_by_relationship: dict[str, list] = {}
+        for edge in edges:
+            relationship_name = edge[2]
+            if edges_by_relationship.get(relationship_name, None):
+                edges_by_relationship[relationship_name].append(edge)
+            else:
+                edges_by_relationship[relationship_name] = [edge]
+
+        results = {}
+        for relationship_name, edges_for_relationship in edges_by_relationship.items():
+            try:
+                # Create the bulk-edge OpenCypher query using UNWIND
+                query = f"""
+                    UNWIND $edges AS edge
+                    MATCH (source:{self._GRAPH_NODE_LABEL})
+                    WHERE id(source) = edge.from_node
+                    MATCH (target:{self._GRAPH_NODE_LABEL})
+                    WHERE id(target) = edge.to_node
+                    MERGE (source)-[r:{relationship_name}]->(target)
+                    ON CREATE SET r = edge.properties, r.updated_at = timestamp()
+                    ON MATCH SET r = edge.properties, r.updated_at = timestamp()
+                    RETURN count(*) AS edges_processed
+                    """
+
+                # Prepare edges data for bulk operation
+                params = {
+                    "edges": [
+                        {
+                            "from_node": str(edge[0]),
+                            "to_node": str(edge[1]),
+                            "relationship_name": relationship_name,
+                            "properties": self._serialize_properties(
+                                edge[3] if len(edge) > 3 and edge[3] else {}
+                            ),
+                        }
+                        for edge in edges_for_relationship
+                    ]
+                }
+                results[relationship_name] = await self.query(query, params)
+            except Exception as e:
+                logger.error(
+                    f"Failed to add edges for relationship {relationship_name}: {format_neptune_error(e)}"
+                )
+                logger.info("Falling back to individual edge creation")
+                for edge in edges_by_relationship:
+                    try:
+                        source_id, target_id, relationship_name = edge[0], edge[1], edge[2]
+                        properties = edge[3] if len(edge) > 3 else {}
+                        await self.add_edge(source_id, target_id, relationship_name, properties)
+                    except Exception as edge_error:
+                        logger.error(
+                            f"Failed to add individual edge {edge[0]} -> {edge[1]}: {format_neptune_error(edge_error)}"
+                        )
+                        continue
+
+        processed_count = 0
+        for result in results.values():
+            processed_count += result[0].get("edges_processed", 0) if result else 0
+        logger.debug(f"Successfully processed {processed_count} edges in bulk operation")
+
+    async def delete_graph(self) -> None:
+        """
+        Delete all nodes and edges from the graph database.
+
+        Returns:
+        --------
+            The result of the query execution, typically indicating success or failure.
+        """
+        try:
+            # Build openCypher query to delete the graph
+            query = f"MATCH (n:{self._GRAPH_NODE_LABEL}) DETACH DELETE n"
+            await self.query(query)
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to delete graph: {error_msg}")
+            raise Exception(f"Failed to delete graph: {error_msg}") from e
+
+    async def get_graph_data(self) -> Tuple[List[Node], List[EdgeData]]:
+        """
+        Retrieve all nodes and edges within the graph.
+
+        Returns:
+        --------
+            - Tuple[List[Node], List[EdgeData]]: A tuple containing all nodes and edges in the graph.
+        """
+        try:
+            # Query to get all nodes
+            nodes_query = f"""
+            MATCH (n:{self._GRAPH_NODE_LABEL})
+            RETURN id(n) AS node_id, properties(n) AS properties
+            """
+
+            # Query to get all edges
+            edges_query = f"""
+            MATCH (source:{self._GRAPH_NODE_LABEL})-[r]->(target:{self._GRAPH_NODE_LABEL})
+            RETURN id(source) AS source_id, id(target) AS target_id, type(r) AS relationship_name, properties(r) AS properties
+            """
+
+            # Execute both queries
+            nodes_result = await self.query(nodes_query)
+            edges_result = await self.query(edges_query)
+
+            # Format nodes as (node_id, properties) tuples
+            nodes = [(result["node_id"], result["properties"]) for result in nodes_result]
+
+            # Format edges as (source_id, target_id, relationship_name, properties) tuples
+            edges = [
+                (
+                    result["source_id"],
+                    result["target_id"],
+                    result["relationship_name"],
+                    result["properties"],
+                )
+                for result in edges_result
+            ]
+
+            logger.debug(f"Retrieved {len(nodes)} nodes and {len(edges)} edges from graph")
+            return (nodes, edges)
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to get graph data: {error_msg}")
+            raise Exception(f"Failed to get graph data: {error_msg}") from e
+
+    async def get_graph_metrics(self, include_optional: bool = False) -> Dict[str, Any]:
+        """
+        Fetch metrics and statistics of the graph, possibly including optional details.
+
+        Parameters:
+        -----------
+            - include_optional (bool): Flag indicating whether to include optional metrics or not.
+
+        Returns:
+        --------
+            - Dict[str, Any]: A dictionary containing graph metrics and statistics.
+        """
+        num_nodes, num_edges = await self._get_model_independent_graph_data()
+        num_cluster, list_clsuter_size = await self._get_connected_components_stat()
+
+        mandatory_metrics = {
+            "num_nodes": num_nodes,
+            "num_edges": num_edges,
+            "mean_degree": (2 * num_edges) / num_nodes if num_nodes != 0 else None,
+            "edge_density": num_edges * 1.0 / (num_nodes * (num_nodes - 1))
+            if num_nodes != 0
+            else None,
+            "num_connected_components": num_cluster,
+            "sizes_of_connected_components": list_clsuter_size,
+        }
+
+        optional_metrics = {
+            "num_selfloops": -1,
+            "diameter": -1,
+            "avg_shortest_path_length": -1,
+            "avg_clustering": -1,
+        }
+
+        if include_optional:
+            optional_metrics["num_selfloops"] = await self._count_self_loops()
+            # Unsupported due to long-running queries when computing the shortest path for each node in the graph:
+            # optional_metrics['diameter']
+            # optional_metrics['avg_shortest_path_length']
+            #
+            # Unsupported due to incompatible algorithm: localClusteringCoefficient
+            # optional_metrics['avg_clustering']
+
+        return mandatory_metrics | optional_metrics
+
+    async def has_edge(self, source_id: str, target_id: str, relationship_name: str) -> bool:
+        """
+        Verify if an edge exists between two specified nodes.
+
+        Parameters:
+        -----------
+            - source_id (str): Unique identifier of the source node.
+            - target_id (str): Unique identifier of the target node.
+            - relationship_name (str): Name of the relationship to verify.
+
+        Returns:
+        --------
+            - bool: True if the edge exists, False otherwise.
+        """
+        try:
+            # Build openCypher query to check if the edge exists
+            query = f"""
+            MATCH (source:{self._GRAPH_NODE_LABEL})-[r:{relationship_name}]->(target:{self._GRAPH_NODE_LABEL})
+            WHERE id(source) = $source_id AND id(target) = $target_id
+            RETURN COUNT(r) > 0 AS edge_exists
+            """
+
+            params = {
+                "source_id": source_id,
+                "target_id": target_id,
+            }
+
+            result = await self.query(query, params)
+
+            if result and len(result) > 0:
+                edge_exists = result.pop().get("edge_exists", False)
+                logger.debug(
+                    f"Edge existence check for "
+                    f"{source_id} -[{relationship_name}]-> {target_id}: {edge_exists}"
+                )
+                return edge_exists
+            else:
+                return False
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to check edge existence {source_id} -> {target_id}: {error_msg}")
+            return False
+
+    async def has_edges(self, edges: List[EdgeData]) -> List[EdgeData]:
+        """
+        Determine the existence of multiple edges in the graph.
+
+        Parameters:
+        -----------
+            - edges (List[EdgeData]): A list of EdgeData objects to check for existence in the graph.
+
+        Returns:
+        --------
+            - List[EdgeData]: A list of EdgeData objects that exist in the graph.
+        """
+        query = f"""
+        UNWIND $edges AS edge
+        MATCH (a:{self._GRAPH_NODE_LABEL})-[r]->(b:{self._GRAPH_NODE_LABEL})
+        WHERE id(a) = edge.from_node AND id(b) = edge.to_node AND type(r) = edge.relationship_name
+        RETURN edge.from_node AS from_node, edge.to_node AS to_node, edge.relationship_name AS relationship_name, count(r) > 0 AS edge_exists
+        """
+
+        try:
+            params = {
+                "edges": [
+                    {
+                        "from_node": str(edge[0]),
+                        "to_node": str(edge[1]),
+                        "relationship_name": edge[2],
+                    }
+                    for edge in edges
+                ],
+            }
+
+            results = await self.query(query, params)
+            logger.debug(f"Found {len(results)} existing edges out of {len(edges)} checked")
+            return [result["edge_exists"] for result in results]
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to check edges existence: {error_msg}")
+            return []
+
+    async def get_edges(self, node_id: str) -> List[EdgeData]:
+        """
+        Retrieve all edges that are connected to the specified node.
+
+        Parameters:
+        -----------
+            - node_id (str): Unique identifier of the node whose edges are to be retrieved.
+
+        Returns:
+        --------
+            - List[EdgeData]: A list of EdgeData objects representing edges connected to the node.
+        """
+        try:
+            # Query to get all edges connected to the node (both incoming and outgoing)
+            query = f"""
+            MATCH (n:{self._GRAPH_NODE_LABEL})-[r]-(m:{self._GRAPH_NODE_LABEL})
+            WHERE id(n) = $node_id
+            RETURN
+                id(n) AS source_id,
+                id(m) AS target_id,
+                type(r) AS relationship_name,
+                properties(r) AS properties
+            """
+
+            params = {"node_id": node_id}
+            result = await self.query(query, params)
+
+            # Format edges as EdgeData tuples: (source_id, target_id, relationship_name, properties)
+            edges = [self._convert_relationship_to_edge(record) for record in result]
+
+            logger.debug(f"Retrieved {len(edges)} edges for node: {node_id}")
+            return edges
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to get edges for node {node_id}: {error_msg}")
+            raise Exception(f"Failed to get edges: {error_msg}") from e
+
+    async def get_disconnected_nodes(self) -> list[str]:
+        """
+        Find and return nodes that are not connected to any other nodes in the graph.
+
+        Returns:
+        --------
+
+            - list[str]: A list of IDs of disconnected nodes.
+        """
+        query = f"""
+            MATCH(n :{self._GRAPH_NODE_LABEL})
+            WHERE NOT (n)--()
+            RETURN COLLECT(ID(n)) as ids
+        """
+
+        results = await self.query(query)
+        return results[0]["ids"] if len(results) > 0 else []
+
+    async def get_predecessors(self, node_id: str, edge_label: str = "") -> list[str]:
+        """
+        Retrieve the predecessor nodes of a specified node based on an optional edge label.
+
+        Parameters:
+        -----------
+
+            - node_id (str): The ID of the node whose predecessors are to be retrieved.
+            - edge_label (str): Optional edge label to filter predecessors. (default None)
+
+        Returns:
+        --------
+
+            - list[str]: A list of predecessor node IDs.
+        """
+
+        edge_label = f" :{edge_label}" if edge_label is not None else ""
+        query = f"""
+        MATCH (node)<-[r{edge_label}]-(predecessor)
+        WHERE node.id = $node_id
+        RETURN predecessor
+        """
+
+        results = await self.query(query, {"node_id": node_id})
+
+        return [result["predecessor"] for result in results]
+
+    async def get_successors(self, node_id: str, edge_label: str = "") -> list[str]:
+        """
+        Retrieve the successor nodes of a specified node based on an optional edge label.
+
+        Parameters:
+        -----------
+
+            - node_id (str): The ID of the node whose successors are to be retrieved.
+            - edge_label (str): Optional edge label to filter successors. (default None)
+
+        Returns:
+        --------
+
+            - list[str]: A list of successor node IDs.
+        """
+
+        edge_label = f" :{edge_label}" if edge_label is not None else ""
+        query = f"""
+        MATCH (node)-[r {edge_label}]->(successor)
+        WHERE node.id = $node_id
+        RETURN successor
+        """
+
+        results = await self.query(query, {"node_id": node_id})
+
+        return [result["successor"] for result in results]
+
+    async def get_neighbors(self, node_id: str) -> List[NodeData]:
+        """
+        Get all neighboring nodes connected to the specified node.
+
+        Parameters:
+        -----------
+            - node_id (str): Unique identifier of the node for which to retrieve neighbors.
+
+        Returns:
+        --------
+            - List[NodeData]: A list of NodeData objects representing neighboring nodes.
+        """
+        try:
+            # Query to get all neighboring nodes (both incoming and outgoing connections)
+            query = f"""
+            MATCH (n:{self._GRAPH_NODE_LABEL})-[r]-(neighbor:{self._GRAPH_NODE_LABEL})
+            WHERE id(n) = $node_id
+            RETURN DISTINCT id(neighbor) AS neighbor_id, properties(neighbor) AS properties
+            """
+
+            params = {"node_id": node_id}
+            result = await self.query(query, params)
+
+            # Format neighbors as NodeData objects
+            neighbors = [
+                {"id": neighbor["neighbor_id"], **neighbor["properties"]} for neighbor in result
+            ]
+
+            logger.debug(f"Retrieved {len(neighbors)} neighbors for node: {node_id}")
+            return neighbors
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to get neighbors for node {node_id}: {error_msg}")
+            raise Exception(f"Failed to get neighbors: {error_msg}") from e
+
+    async def get_nodeset_subgraph(
+        self, node_type: Type[Any], node_name: List[str]
+    ) -> Tuple[List[Tuple[int, dict]], List[Tuple[int, int, str, dict]]]:
+        """
+        Fetch a subgraph consisting of a specific set of nodes and their relationships.
+
+        Parameters:
+        -----------
+            - node_type (Type[Any]): The type of nodes to include in the subgraph.
+            - node_name (List[str]): A list of names of the nodes to include in the subgraph.
+
+        Returns:
+        --------
+            - Tuple[List[Tuple[int, dict]], List[Tuple[int, int, str, dict]]]: A tuple containing nodes and edges of the subgraph.
+        """
+        try:
+            # Query to get nodes by name and their connected subgraph
+            query = f"""
+            UNWIND $names AS wantedName
+            MATCH (n:{self._GRAPH_NODE_LABEL})
+            WHERE n.name = wantedName AND n.type = $type
+            WITH collect(DISTINCT n) AS primary
+            UNWIND primary AS p
+            OPTIONAL MATCH (p)-[r]-(nbr:{self._GRAPH_NODE_LABEL})
+            WITH primary, collect(DISTINCT nbr) AS nbrs, collect(DISTINCT r) AS rels
+            WITH primary + nbrs AS nodelist, rels
+            UNWIND nodelist AS node
+            WITH collect(DISTINCT node) AS nodes, rels
+            MATCH (a:{self._GRAPH_NODE_LABEL})-[r]-(b:{self._GRAPH_NODE_LABEL})
+            WHERE a IN nodes AND b IN nodes
+            WITH nodes, collect(DISTINCT r) AS all_rels
+            RETURN
+              [n IN nodes | {{
+                id: id(n),
+                properties: properties(n)
+              }}] AS rawNodes,
+              [r IN all_rels | {{
+                source_id: id(startNode(r)),
+                target_id: id(endNode(r)),
+                type: type(r),
+                properties: properties(r)
+              }}] AS rawRels
+            """
+
+            params = {"names": node_name, "type": node_type.__name__}
+
+            result = await self.query(query, params)
+
+            if not result:
+                logger.debug(f"No subgraph found for node type {node_type} with names {node_name}")
+                return ([], [])
+
+            raw_nodes = result[0]["rawNodes"]
+            raw_rels = result[0]["rawRels"]
+
+            # Format nodes as (node_id, properties) tuples
+            nodes = [(n["id"], n["properties"]) for n in raw_nodes]
+
+            # Format edges as (source_id, target_id, relationship_name, properties) tuples
+            edges = [(r["source_id"], r["target_id"], r["type"], r["properties"]) for r in raw_rels]
+
+            logger.debug(
+                f"Retrieved subgraph with {len(nodes)} nodes and {len(edges)} edges for type {node_type.__name__}"
+            )
+            return (nodes, edges)
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to get nodeset subgraph for type {node_type}: {error_msg}")
+            raise Exception(f"Failed to get nodeset subgraph: {error_msg}") from e
+
+    async def get_connections(self, node_id: UUID) -> list:
+        """
+        Get all nodes connected to a specified node and their relationship details.
+
+        Parameters:
+        -----------
+            - node_id (str): Unique identifier of the node for which to retrieve connections.
+
+        Returns:
+        --------
+            - List[Tuple[NodeData, Dict[str, Any], NodeData]]: A list of tuples containing connected nodes and relationship details.
+        """
+        try:
+            # Query to get all connections (both incoming and outgoing)
+            query = f"""
+            MATCH (source:{self._GRAPH_NODE_LABEL})-[r]->(target:{self._GRAPH_NODE_LABEL})
+            WHERE id(source) = $node_id OR id(target) = $node_id
+            RETURN
+                id(source) AS source_id,
+                properties(source) AS source_props,
+                id(target) AS target_id,
+                properties(target) AS target_props,
+                type(r) AS relationship_name,
+                properties(r) AS relationship_props
+            """
+
+            params = {"node_id": str(node_id)}
+            result = await self.query(query, params)
+
+            connections = []
+            for record in result:
+                # Return as (source_node, relationship, target_node)
+                connections.append(
+                    (
+                        {"id": record["source_id"], **record["source_props"]},
+                        {
+                            "relationship_name": record["relationship_name"],
+                            **record["relationship_props"],
+                        },
+                        {"id": record["target_id"], **record["target_props"]},
+                    )
+                )
+
+            logger.debug(f"Retrieved {len(connections)} connections for node: {node_id}")
+            return connections
+
+        except Exception as e:
+            error_msg = format_neptune_error(e)
+            logger.error(f"Failed to get connections for node {node_id}: {error_msg}")
+            raise Exception(f"Failed to get connections: {error_msg}") from e
+
+    async def remove_connection_to_predecessors_of(self, node_ids: list[str], edge_label: str):
+        """
+        Remove connections (edges) to all predecessors of specified nodes based on edge label.
+
+        Parameters:
+        -----------
+
+            - node_ids (list[str]): A list of IDs of nodes from which connections are to be
+              removed.
+            - edge_label (str): The label of the edges to remove.
+
+        """
+        query = f"""
+        UNWIND $node_ids AS node_id
+        MATCH ({{`~id`: node_id}})-[r:{edge_label}]->(predecessor)
+        DELETE r;
+        """
+        params = {"node_ids": node_ids}
+        await self.query(query, params)
+
+    async def remove_connection_to_successors_of(self, node_ids: list[str], edge_label: str):
+        """
+        Remove connections (edges) to all successors of specified nodes based on edge label.
+
+        Parameters:
+        -----------
+
+            - node_ids (list[str]): A list of IDs of nodes from which connections are to be
+              removed.
+            - edge_label (str): The label of the edges to remove.
+
+        """
+        query = f"""
+        UNWIND $node_ids AS node_id
+        MATCH ({{`~id`: node_id}})<-[r:{edge_label}]-(successor)
+        DELETE r;
+        """
+        params = {"node_ids": node_ids}
+        await self.query(query, params)
+
+    async def get_node_labels_string(self):
+        """
+        Fetch all node labels from the database and return them as a formatted string.
+
+        Returns:
+        --------
+
+            A formatted string of node labels.
+
+        Raises:
+        -------
+            ValueError: If no node labels are found in the database.
+        """
+        node_labels_query = (
+            "CALL neptune.graph.pg_schema() YIELD schema RETURN schema.nodeLabels as labels "
+        )
+        node_labels_result = await self.query(node_labels_query)
+        node_labels = node_labels_result[0]["labels"] if node_labels_result else []
+
+        if not node_labels:
+            raise ValueError("No node labels found in the database")
+
+        return str(node_labels)
+
+    async def get_relationship_labels_string(self):
+        """
+        Fetch all relationship types from the database and return them as a formatted string.
+
+        Returns:
+        --------
+
+            A formatted string of relationship types.
+        """
+        relationship_types_query = (
+            "CALL neptune.graph.pg_schema() YIELD schema RETURN schema.edgeLabels as relationships "
+        )
+        relationship_types_result = await self.query(relationship_types_query)
+        relationship_types = (
+            relationship_types_result[0]["relationships"] if relationship_types_result else []
+        )
+
+        if not relationship_types:
+            raise ValueError("No relationship types found in the database.")
+
+        relationship_types_undirected_str = (
+            "{"
+            + ", ".join(f"{rel}" + ": {orientation: 'UNDIRECTED'}" for rel in relationship_types)
+            + "}"
+        )
+        return relationship_types_undirected_str
+
+    async def drop_graph(self, graph_name="myGraph"):
+        """
+        Drop an existing graph from the database based on its name.
+
+        Note: This method is currently a placeholder because GDS (Graph Data Science)
+        projection is not supported in Neptune Analytics.
+
+        Parameters:
+        -----------
+
+            - graph_name: The name of the graph to drop, defaults to 'myGraph'. (default
+              'myGraph')
+        """
+        pass
+
+    async def graph_exists(self, graph_name="myGraph"):
+        """
+        Check if a graph with a given name exists in the database.
+
+        Note: This method is currently a placeholder because GDS (Graph Data Science)
+        projection is not supported in Neptune Analytics.
+
+        Parameters:
+        -----------
+
+            - graph_name: The name of the graph to check for existence, defaults to 'myGraph'.
+              (default 'myGraph')
+
+        Returns:
+        --------
+
+            True if the graph exists, otherwise False.
+        """
+        pass
+
+    async def project_entire_graph(self, graph_name="myGraph"):
+        """
+        Project all node labels and relationship types into an in-memory graph using GDS.
+
+        Note: This method is currently a placeholder because GDS (Graph Data Science)
+        projection is not supported in Neptune Anlaytics.
+        """
+        pass
+
+    async def get_filtered_graph_data(self, attribute_filters: list[dict[str, list]]):
+        """
+        Fetch nodes and edges filtered by specific attribute criteria.
+
+        Parameters:
+        -----------
+
+            - attribute_filters: A list of dictionaries representing attributes and associated
+              values for filtering.
+
+        Returns:
+        --------
+
+            A tuple containing filtered nodes and edges based on the specified criteria.
+        """
+        where_clauses_n = []
+        where_clauses_m = []
+        for attribute, values in attribute_filters[0].items():
+            values_str = ", ".join(
+                f"'{value}'" if isinstance(value, str) else str(value) for value in values
+            )
+            where_clauses_n.append(f"n.{attribute} IN [{values_str}]")
+            where_clauses_m.append(f"m.{attribute} IN [{values_str}]")
+
+        node_where_clauses_n_str = " AND ".join(where_clauses_n)
+        node_where_clauses_m_str = " AND ".join(where_clauses_m)
+        edge_where_clause = f"{node_where_clauses_n_str} AND {node_where_clauses_m_str}"
+
+        query_nodes = f"""
+           MATCH (n :{self._GRAPH_NODE_LABEL})
+           WHERE {node_where_clauses_n_str}
+           RETURN ID(n) AS id, labels(n) AS labels, properties(n) AS properties
+           """
+        result_nodes = await self.query(query_nodes)
+
+        nodes = [
+            (
+                record["id"],
+                record["properties"],
+            )
+            for record in result_nodes
+        ]
+
+        query_edges = f"""
+           MATCH (n :{self._GRAPH_NODE_LABEL})-[r]->(m :{self._GRAPH_NODE_LABEL})
+           WHERE {edge_where_clause}
+           RETURN ID(n) AS source, ID(m) AS target, TYPE(r) AS type, properties(r) AS properties
+           """
+        result_edges = await self.query(query_edges)
+
+        edges = [
+            (
+                record["source"],
+                record["target"],
+                record["type"],
+                record["properties"],
+            )
+            for record in result_edges
+        ]
+
+        return (nodes, edges)
+
+    async def get_degree_one_nodes(self, node_type: str):
+        """
+        Fetch nodes of a specified type that have exactly one connection.
+
+        Parameters:
+        -----------
+
+            - node_type (str): The type of nodes to retrieve, must be 'Entity' or 'EntityType'.
+
+        Returns:
+        --------
+
+            A list of nodes with exactly one connection of the specified type.
+        """
+        if not node_type or node_type not in ["Entity", "EntityType"]:
+            raise ValueError("node_type must be either 'Entity' or 'EntityType'")
+
+        query = f"""
+                MATCH (n :{self._GRAPH_NODE_LABEL})
+                WHERE size((n)--()) = 1
+                AND n.type = $node_type
+                RETURN n
+                """
+        result = await self.query(query, {"node_type": node_type})
+        return [record["n"] for record in result] if result else []
+
+    async def get_document_subgraph(self, data_id: str):
+        """
+        Retrieve a subgraph related to a document identified by its content hash, including
+        related entities and chunks.
+
+        Parameters:
+        -----------
+
+            - data_id (str): The document_id identifying the document whose subgraph should be
+              retrieved.
+
+        Returns:
+        --------
+
+            The subgraph data as a dictionary, or None if not found.
+        """
+        query = f"""
+
+        MATCH (doc)
+        WHERE (doc:{self._GRAPH_NODE_LABEL})
+        AND doc.type in ['TextDocument', 'PdfDocument']
+        AND doc.id = $data_id
+
+        OPTIONAL MATCH (doc)<-[:is_part_of]-(chunk {{type: 'DocumentChunk'}})
+
+        // Alternative to WHERE NOT EXISTS
+        OPTIONAL MATCH (chunk)-[:contains]->(entity {{type: 'Entity'}})
+        OPTIONAL MATCH (entity)<-[:contains]-(otherChunk {{type: 'DocumentChunk'}})-[:is_part_of]->(otherDoc)
+          WHERE otherDoc.type in ['TextDocument', 'PdfDocument']
+          AND otherDoc.id <> doc.id
+                OPTIONAL MATCH (chunk)<-[:made_from]-(made_node {{type: 'TextSummary'}})
+
+        OPTIONAL MATCH (chunk)<-[:made_from]-(made_node {{type: 'TextSummary'}})
+
+        // Alternative to WHERE NOT EXISTS
+        OPTIONAL MATCH (entity)-[:is_a]->(type {{type: 'EntityType'}})
+        OPTIONAL MATCH (type)<-[:is_a]-(otherEntity {{type: 'Entity'}})<-[:contains]-(otherChunk {{type: 'DocumentChunk'}})-[:is_part_of]->(otherDoc)
+          WHERE otherDoc.type in ['TextDocument', 'PdfDocument']
+          AND otherDoc.id <> doc.id
+
+        // Alternative to WHERE NOT EXISTS
+        WITH doc, entity, chunk, made_node, type, otherDoc
+        WHERE otherDoc IS NULL
+
+        RETURN
+            collect(DISTINCT doc) as document,
+            collect(DISTINCT chunk) as chunks,
+            collect(DISTINCT entity) as orphan_entities,
+            collect(DISTINCT made_node) as made_from_nodes,
+            collect(DISTINCT type) as orphan_types
+        """
+        result = await self.query(query, {"data_id": data_id})
+        return result[0] if result else None
+
+    async def _get_model_independent_graph_data(self):
+        """
+        Retrieve the basic graph data without considering the model specifics, returning nodes
+        and edges.
+
+        Returns:
+        --------
+
+            A tuple of nodes and edges data.
+        """
+        query_string = f"""
+            MATCH (n :{self._GRAPH_NODE_LABEL})
+            WITH count(n) AS nodeCount
+            MATCH (a :{self._GRAPH_NODE_LABEL})-[r]->(b :{self._GRAPH_NODE_LABEL})
+            RETURN nodeCount AS numVertices, count(r) AS numEdges
+        """
+        query_response = await self.query(query_string)
+        num_nodes = query_response[0].get("numVertices")
+        num_edges = query_response[0].get("numEdges")
+
+        return (num_nodes, num_edges)
+
+    async def _get_connected_components_stat(self):
+        """
+        Retrieve statistics about connected components in the graph.
+
+        This method analyzes the graph to find all connected components
+        and returns both the sizes of each component and the total number of components.
+
+
+        Returns:
+        --------
+            tuple[list[int], int]
+            A tuple containing:
+              - A list of sizes for each connected component (descending order).
+              - The total number of connected components.
+            Returns ([], 0) if no connected components are found.
+        """
+        query = f"""
+        MATCH(n :{self._GRAPH_NODE_LABEL})
+        CALL neptune.algo.wcc(n,{{}})
+        YIELD node, component
+        RETURN component, count(*) AS size
+        ORDER BY size DESC
+        """
+
+        result = await self.query(query)
+        size_connected_components = [record["size"] for record in result] if result else []
+        num_connected_components = len(result)
+
+        return (size_connected_components, num_connected_components)
+
+    async def _count_self_loops(self):
+        """
+        Count the number of self-loop relationships in the Neptune Anlaytics graph backend.
+
+        This function executes a OpenCypher query to find and count all edge relationships that
+        begin and end at the same node (self-loops). It returns the count of such relationships
+        or 0 if no results are found.
+
+        Returns:
+        --------
+
+            The count of self-loop relationships found in the database, or 0 if none were found.
+        """
+        query = f"""
+        MATCH (n :{self._GRAPH_NODE_LABEL})-[r]->(n :{self._GRAPH_NODE_LABEL})
+        RETURN count(r) AS adapter_loop_count;
+        """
+        result = await self.query(query)
+        return result[0]["adapter_loop_count"] if result else 0
+
+    @staticmethod
+    def _convert_relationship_to_edge(relationship: dict) -> EdgeData:
+        return (
+            relationship["source_id"],
+            relationship["target_id"],
+            relationship["relationship_name"],
+            relationship["properties"],
+        )