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

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

@@ -0,0 +1,115 @@
+"""Neptune Analytics Exceptions
+
+This module defines custom exceptions for Neptune Analytics operations.
+"""
+
+from cognee.exceptions import CogneeApiError
+from fastapi import status
+
+
+class NeptuneAnalyticsError(CogneeApiError):
+    """Base exception for Neptune Analytics operations."""
+
+    def __init__(
+        self,
+        message: str = "Neptune Analytics error.",
+        name: str = "NeptuneAnalyticsError",
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+    ):
+        super().__init__(message, name, status_code)
+
+
+class NeptuneAnalyticsConnectionError(NeptuneAnalyticsError):
+    """Exception raised when connection to Neptune Analytics fails."""
+
+    def __init__(
+        self,
+        message: str = "Unable to connect to Neptune Analytics. Please check the endpoint and network connectivity.",
+        name: str = "NeptuneAnalyticsConnectionError",
+        status_code=status.HTTP_404_NOT_FOUND,
+    ):
+        super().__init__(message, name, status_code)
+
+
+class NeptuneAnalyticsQueryError(NeptuneAnalyticsError):
+    """Exception raised when a query execution fails."""
+
+    def __init__(
+        self,
+        message: str = "The query execution failed due to invalid syntax or semantic issues.",
+        name: str = "NeptuneAnalyticsQueryError",
+        status_code=status.HTTP_400_BAD_REQUEST,
+    ):
+        super().__init__(message, name, status_code)
+
+
+class NeptuneAnalyticsAuthenticationError(NeptuneAnalyticsError):
+    """Exception raised when authentication with Neptune Analytics fails."""
+
+    def __init__(
+        self,
+        message: str = "Authentication with Neptune Analytics failed. Please verify your credentials.",
+        name: str = "NeptuneAnalyticsAuthenticationError",
+        status_code=status.HTTP_401_UNAUTHORIZED,
+    ):
+        super().__init__(message, name, status_code)
+
+
+class NeptuneAnalyticsConfigurationError(NeptuneAnalyticsError):
+    """Exception raised when Neptune Analytics configuration is invalid."""
+
+    def __init__(
+        self,
+        message: str = "Neptune Analytics configuration is invalid or incomplete. Please review your setup.",
+        name: str = "NeptuneAnalyticsConfigurationError",
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+    ):
+        super().__init__(message, name, status_code)
+
+
+class NeptuneAnalyticsTimeoutError(NeptuneAnalyticsError):
+    """Exception raised when a Neptune Analytics operation times out."""
+
+    def __init__(
+        self,
+        message: str = "The operation timed out while communicating with Neptune Analytics.",
+        name: str = "NeptuneAnalyticsTimeoutError",
+        status_code=status.HTTP_504_GATEWAY_TIMEOUT,
+    ):
+        super().__init__(message, name, status_code)
+
+
+class NeptuneAnalyticsThrottlingError(NeptuneAnalyticsError):
+    """Exception raised when requests are throttled by Neptune Analytics."""
+
+    def __init__(
+        self,
+        message: str = "Request was throttled by Neptune Analytics due to exceeding rate limits.",
+        name: str = "NeptuneAnalyticsThrottlingError",
+        status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+    ):
+        super().__init__(message, name, status_code)
+
+
+class NeptuneAnalyticsResourceNotFoundError(NeptuneAnalyticsError):
+    """Exception raised when a Neptune Analytics resource is not found."""
+
+    def __init__(
+        self,
+        message: str = "The requested Neptune Analytics resource could not be found.",
+        name: str = "NeptuneAnalyticsResourceNotFoundError",
+        status_code=status.HTTP_404_NOT_FOUND,
+    ):
+        super().__init__(message, name, status_code)
+
+
+class NeptuneAnalyticsInvalidParameterError(NeptuneAnalyticsError):
+    """Exception raised when invalid parameters are provided to Neptune Analytics."""
+
+    def __init__(
+        self,
+        message: str = "One or more parameters provided to Neptune Analytics are invalid or missing.",
+        name: str = "NeptuneAnalyticsInvalidParameterError",
+        status_code=status.HTTP_400_BAD_REQUEST,
+    ):
+        super().__init__(message, name, status_code)

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

@@ -0,0 +1,224 @@
+"""Neptune Utilities
+
+This module provides utility functions for Neptune Analytics operations including
+connection management, URL parsing, and Neptune-specific configurations.
+"""
+
+import re
+from typing import Optional, Dict, Any, Tuple
+from urllib.parse import urlparse
+
+from cognee.shared.logging_utils import get_logger
+
+logger = get_logger("NeptuneUtils")
+
+
+def parse_neptune_url(url: str) -> Tuple[str, str]:
+    """
+    Parse a Neptune Analytics URL to extract graph ID and region.
+
+    Expected format: neptune-graph://<GRAPH_ID>?region=<REGION>
+    or neptune-graph://<GRAPH_ID> (defaults to us-east-1)
+
+    Parameters:
+    -----------
+        - url (str): The Neptune Analytics URL to parse
+
+    Returns:
+    --------
+        - Tuple[str, str]: A tuple containing (graph_id, region)
+
+    Raises:
+    -------
+        - ValueError: If the URL format is invalid
+    """
+    try:
+        parsed = urlparse(url)
+
+        if parsed.scheme != "neptune-graph":
+            raise ValueError(f"Invalid scheme: {parsed.scheme}. Expected 'neptune-graph'")
+
+        graph_id = parsed.hostname or parsed.path.lstrip("/")
+        if not graph_id:
+            raise ValueError("Graph ID not found in URL")
+
+        # Extract region from query parameters
+        region = "us-east-1"  # default region
+        if parsed.query:
+            query_params = dict(
+                param.split("=") for param in parsed.query.split("&") if "=" in param
+            )
+            region = query_params.get("region", region)
+
+        return graph_id, region
+
+    except Exception as e:
+        raise ValueError(f"Failed to parse Neptune Analytics URL '{url}': {str(e)}")
+
+
+def validate_graph_id(graph_id: str) -> bool:
+    """
+    Validate a Neptune Analytics graph ID format.
+
+    Graph IDs should follow AWS naming conventions.
+
+    Parameters:
+    -----------
+        - graph_id (str): The graph ID to validate
+
+    Returns:
+    --------
+        - bool: True if the graph ID is valid, False otherwise
+    """
+    if not graph_id:
+        return False
+
+    # Neptune Analytics graph IDs should be alphanumeric with hyphens
+    # and between 1-63 characters
+    pattern = r"^[a-zA-Z0-9][a-zA-Z0-9\-]{0,62}$"
+    return bool(re.match(pattern, graph_id))
+
+
+def validate_aws_region(region: str) -> bool:
+    """
+    Validate an AWS region format.
+
+    Parameters:
+    -----------
+        - region (str): The AWS region to validate
+
+    Returns:
+    --------
+        - bool: True if the region format is valid, False otherwise
+    """
+    if not region:
+        return False
+
+    # AWS regions follow the pattern: us-east-1, eu-west-1, etc.
+    pattern = r"^[a-z]{2,3}-[a-z]+-\d+$"
+    return bool(re.match(pattern, region))
+
+
+def build_neptune_config(
+    graph_id: str,
+    region: Optional[str],
+    aws_access_key_id: Optional[str] = None,
+    aws_secret_access_key: Optional[str] = None,
+    aws_session_token: Optional[str] = None,
+    **kwargs,
+) -> Dict[str, Any]:
+    """
+    Build a configuration dictionary for Neptune Analytics connection.
+
+    Parameters:
+    -----------
+        - graph_id (str): The Neptune Analytics graph identifier
+        - region (Optional[str]): AWS region where the graph is located
+        - 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
+        - **kwargs: Additional configuration parameters
+
+    Returns:
+    --------
+        - Dict[str, Any]: Configuration dictionary for Neptune Analytics
+
+    Raises:
+    -------
+        - ValueError: If required parameters are invalid
+    """
+    config = {
+        "graph_id": graph_id,
+        "service_name": "neptune-graph",
+    }
+
+    # Add AWS credentials if provided
+    if region:
+        config["region"] = region
+
+    if aws_access_key_id:
+        config["aws_access_key_id"] = aws_access_key_id
+
+    if aws_secret_access_key:
+        config["aws_secret_access_key"] = aws_secret_access_key
+
+    if aws_session_token:
+        config["aws_session_token"] = aws_session_token
+
+    # Add any additional configuration
+    config.update(kwargs)
+
+    return config
+
+
+def get_neptune_endpoint_url(graph_id: str, region: str) -> str:
+    """
+    Construct the Neptune Analytics endpoint URL for a given graph and region.
+
+    Parameters:
+    -----------
+        - graph_id (str): The Neptune Analytics graph identifier
+        - region (str): AWS region where the graph is located
+
+    Returns:
+    --------
+        - str: The Neptune Analytics endpoint URL
+    """
+    return f"https://neptune-graph.{region}.amazonaws.com/graphs/{graph_id}"
+
+
+def format_neptune_error(error: Exception) -> str:
+    """
+    Format Neptune Analytics specific errors for better readability.
+
+    Parameters:
+    -----------
+        - error (Exception): The exception to format
+
+    Returns:
+    --------
+        - str: Formatted error message
+    """
+    error_msg = str(error)
+
+    # Common Neptune Analytics error patterns and their user-friendly messages
+    error_mappings = {
+        "AccessDenied": "Access denied. Please check your AWS credentials and permissions.",
+        "GraphNotFound": "Graph not found. Please verify the graph ID and region.",
+        "InvalidParameter": "Invalid parameter provided. Please check your request parameters.",
+        "ThrottlingException": "Request was throttled. Please retry with exponential backoff.",
+        "InternalServerError": "Internal server error occurred. Please try again later.",
+    }
+
+    for error_type, friendly_msg in error_mappings.items():
+        if error_type in error_msg:
+            return f"{friendly_msg} Original error: {error_msg}"
+
+    return error_msg
+
+
+def get_default_query_timeout() -> int:
+    """
+    Get the default query timeout for Neptune Analytics operations.
+
+    Returns:
+    --------
+        - int: Default timeout in seconds
+    """
+    return 300  # 5 minutes
+
+
+def get_default_connection_config() -> Dict[str, Any]:
+    """
+    Get default connection configuration for Neptune Analytics.
+
+    Returns:
+    --------
+        - Dict[str, Any]: Default connection configuration
+    """
+    return {
+        "query_timeout": get_default_query_timeout(),
+        "max_retries": 3,
+        "retry_delay": 1.0,
+        "preferred_query_language": "openCypher",
+    }

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

@@ -826,7 +826,7 @@ class NetworkXAdapter(GraphDBInterface):
 
         return mandatory_metrics | optional_metrics
 
-    async def get_document_subgraph(self, content_hash: str):
+    async def get_document_subgraph(self, data_id: str):
         """
         Retrieve all relevant nodes when a document is being deleted, including chunks and
         orphaned entities.
@@ -834,7 +834,7 @@ class NetworkXAdapter(GraphDBInterface):
         Parameters:
         -----------
 
-            - content_hash (str): The hash identifying the content of the document to fetch
+            - data_id(str): The data id identifying the document to fetch
               related nodes for.
 
         Returns:
@@ -853,7 +853,7 @@ class NetworkXAdapter(GraphDBInterface):
         for node_id, attrs in self.graph.nodes(data=True):
             if (
                 attrs.get("type") in ["TextDocument", "PdfDocument"]
-                and attrs.get("name") == f"text_{content_hash}"
+                and attrs.get("id") == f"{data_id}"
             ):
                 document = {"id": str(node_id), **attrs}  # Convert UUID to string for consistency
                 document_node_id = node_id  # Keep the original UUID