cognee 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

cognee/infrastructure/databases/dataset_database_handler/use_dataset_database_handler.py

@@ -0,0 +1,10 @@
+from .supported_dataset_database_handlers import supported_dataset_database_handlers
+
+
+def use_dataset_database_handler(
+    dataset_database_handler_name, dataset_database_handler, dataset_database_provider
+):
+    supported_dataset_database_handlers[dataset_database_handler_name] = {
+        "handler_instance": dataset_database_handler,
+        "handler_provider": dataset_database_provider,
+    }

cognee/infrastructure/databases/exceptions/exceptions.py

@@ -148,3 +148,19 @@ class CacheConnectionError(CogneeConfigurationError):
         status_code: int = status.HTTP_503_SERVICE_UNAVAILABLE,
     ):
         super().__init__(message, name, status_code)
+
+
+class SharedKuzuLockRequiresRedisError(CogneeConfigurationError):
+    """
+    Raised when shared Kuzu locking is requested without configuring the Redis backend.
+    """
+
+    def __init__(
+        self,
+        message: str = (
+            "Shared Kuzu lock requires Redis cache backend. Configure Redis to enable shared Kuzu locking."
+        ),
+        name: str = "SharedKuzuLockRequiresRedisError",
+        status_code: int = status.HTTP_400_BAD_REQUEST,
+    ):
+        super().__init__(message, name, status_code)

cognee/infrastructure/databases/graph/config.py

@@ -26,6 +26,7 @@ class GraphConfig(BaseSettings):
     - graph_database_username
     - graph_database_password
     - graph_database_port
+    - graph_database_key
     - graph_file_path
     - graph_model
     - graph_topology
@@ -41,10 +42,12 @@ class GraphConfig(BaseSettings):
     graph_database_username: str = ""
     graph_database_password: str = ""
     graph_database_port: int = 123
+    graph_database_key: str = ""
     graph_file_path: str = ""
     graph_filename: str = ""
     graph_model: object = KnowledgeGraph
     graph_topology: object = KnowledgeGraph
+    graph_dataset_database_handler: str = "kuzu"
     model_config = SettingsConfigDict(env_file=".env", extra="allow", populate_by_name=True)
 
     # Model validator updates graph_filename and path dynamically after class creation based on current database provider
@@ -90,10 +93,12 @@ class GraphConfig(BaseSettings):
             "graph_database_username": self.graph_database_username,
             "graph_database_password": self.graph_database_password,
             "graph_database_port": self.graph_database_port,
+            "graph_database_key": self.graph_database_key,
             "graph_file_path": self.graph_file_path,
             "graph_model": self.graph_model,
             "graph_topology": self.graph_topology,
             "model_config": self.model_config,
+            "graph_dataset_database_handler": self.graph_dataset_database_handler,
         }
 
     def to_hashable_dict(self) -> dict:
@@ -116,7 +121,9 @@ class GraphConfig(BaseSettings):
             "graph_database_username": self.graph_database_username,
             "graph_database_password": self.graph_database_password,
             "graph_database_port": self.graph_database_port,
+            "graph_database_key": self.graph_database_key,
             "graph_file_path": self.graph_file_path,
+            "graph_dataset_database_handler": self.graph_dataset_database_handler,
         }
 
 

cognee/infrastructure/databases/graph/get_graph_engine.py

@@ -33,6 +33,8 @@ def create_graph_engine(
     graph_database_username="",
     graph_database_password="",
     graph_database_port="",
+    graph_database_key="",
+    graph_dataset_database_handler="",
 ):
     """
     Create a graph engine based on the specified provider type.
@@ -69,6 +71,7 @@ def create_graph_engine(
             graph_database_url=graph_database_url,
             graph_database_username=graph_database_username,
             graph_database_password=graph_database_password,
+            database_name=graph_database_name,
         )
 
     if graph_database_provider == "neo4j":

cognee/infrastructure/databases/graph/graph_db_interface.py

@@ -398,3 +398,18 @@ class GraphDBInterface(ABC):
             - node_id (Union[str, UUID]): Unique identifier of the node for which to retrieve connections.
         """
         raise NotImplementedError
+
+    @abstractmethod
+    async def get_filtered_graph_data(
+        self, attribute_filters: List[Dict[str, List[Union[str, int]]]]
+    ) -> Tuple[List[Node], List[EdgeData]]:
+        """
+        Retrieve nodes and edges filtered by the provided attribute criteria.
+
+        Parameters:
+        -----------
+
+            - attribute_filters: A list of dictionaries where keys are attribute names and values
+              are lists of attribute values to filter by.
+        """
+        raise NotImplementedError

cognee/infrastructure/databases/graph/kuzu/KuzuDatasetDatabaseHandler.py

@@ -0,0 +1,81 @@
+import os
+from uuid import UUID
+from typing import Optional
+
+from cognee.infrastructure.databases.graph.get_graph_engine import create_graph_engine
+from cognee.base_config import get_base_config
+from cognee.modules.users.models import User
+from cognee.modules.users.models import DatasetDatabase
+from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface
+
+
+class KuzuDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
+    """
+    Handler for interacting with Kuzu Dataset databases.
+    """
+
+    @classmethod
+    async def create_dataset(cls, dataset_id: Optional[UUID], user: Optional[User]) -> dict:
+        """
+        Create a new Kuzu instance for the dataset. Return connection info that will be mapped to the dataset.
+
+        Args:
+            dataset_id: Dataset UUID
+            user: User object who owns the dataset and is making the request
+
+        Returns:
+            dict: Connection details for the created Kuzu instance
+
+        """
+        from cognee.infrastructure.databases.graph.config import get_graph_config
+
+        graph_config = get_graph_config()
+
+        if graph_config.graph_database_provider != "kuzu":
+            raise ValueError(
+                "KuzuDatasetDatabaseHandler can only be used with Kuzu graph database provider."
+            )
+
+        graph_db_name = f"{dataset_id}.pkl"
+        graph_db_url = graph_config.graph_database_url
+        graph_db_key = graph_config.graph_database_key
+        graph_db_username = graph_config.graph_database_username
+        graph_db_password = graph_config.graph_database_password
+
+        return {
+            "graph_database_name": graph_db_name,
+            "graph_database_url": graph_db_url,
+            "graph_database_provider": graph_config.graph_database_provider,
+            "graph_database_key": graph_db_key,
+            "graph_dataset_database_handler": "kuzu",
+            "graph_database_connection_info": {
+                "graph_database_username": graph_db_username,
+                "graph_database_password": graph_db_password,
+            },
+        }
+
+    @classmethod
+    async def delete_dataset(cls, dataset_database: DatasetDatabase):
+        base_config = get_base_config()
+        databases_directory_path = os.path.join(
+            base_config.system_root_directory, "databases", str(dataset_database.owner_id)
+        )
+        graph_file_path = os.path.join(
+            databases_directory_path, dataset_database.graph_database_name
+        )
+        graph_engine = create_graph_engine(
+            graph_database_provider=dataset_database.graph_database_provider,
+            graph_database_url=dataset_database.graph_database_url,
+            graph_database_name=dataset_database.graph_database_name,
+            graph_database_key=dataset_database.graph_database_key,
+            graph_file_path=graph_file_path,
+            graph_database_username=dataset_database.graph_database_connection_info.get(
+                "graph_database_username", ""
+            ),
+            graph_database_password=dataset_database.graph_database_connection_info.get(
+                "graph_database_password", ""
+            ),
+            graph_dataset_database_handler="",
+            graph_database_port="",
+        )
+        await graph_engine.delete_graph()

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

@@ -12,6 +12,7 @@ from contextlib import asynccontextmanager
 from concurrent.futures import ThreadPoolExecutor
 from typing import Dict, Any, List, Union, Optional, Tuple, Type
 
+from cognee.exceptions import CogneeValidationError
 from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.utils.run_sync import run_sync
 from cognee.infrastructure.files.storage import get_file_storage
@@ -1186,6 +1187,11 @@ class KuzuAdapter(GraphDBInterface):
               A tuple with two elements: a list of tuples of (node_id, properties) and a list of
               tuples of (source_id, target_id, relationship_name, properties).
         """
+
+        import time
+
+        start_time = time.time()
+
         try:
             nodes_query = """
             MATCH (n:Node)
@@ -1249,6 +1255,11 @@ class KuzuAdapter(GraphDBInterface):
                             },
                         )
                     )
+
+            retrieval_time = time.time() - start_time
+            logger.info(
+                f"Retrieved {len(nodes)} nodes and {len(edges)} edges in {retrieval_time:.2f} seconds"
+            )
             return formatted_nodes, formatted_edges
         except Exception as e:
             logger.error(f"Failed to get graph data: {e}")
@@ -1417,6 +1428,92 @@ class KuzuAdapter(GraphDBInterface):
                 formatted_edges.append((source_id, target_id, rel_type, props))
         return formatted_nodes, formatted_edges
 
+    async def get_id_filtered_graph_data(self, target_ids: list[str]):
+        """
+        Retrieve graph data filtered by specific node IDs, including their direct neighbors
+        and only edges where one endpoint matches those IDs.
+
+        Returns:
+            nodes: List[dict]   -> Each dict includes "id" and all node properties
+            edges: List[dict]   -> Each dict includes "source", "target", "type", "properties"
+        """
+        import time
+
+        start_time = time.time()
+
+        try:
+            if not target_ids:
+                logger.warning("No target IDs provided for ID-filtered graph retrieval.")
+                return [], []
+
+            if not all(isinstance(x, str) for x in target_ids):
+                raise CogneeValidationError("target_ids must be a list of strings")
+
+            query = """
+            MATCH (n:Node)-[r]->(m:Node)
+            WHERE n.id IN $target_ids OR m.id IN $target_ids
+            RETURN n.id, {
+                name: n.name,
+                type: n.type,
+                properties: n.properties
+            }, m.id, {
+                name: m.name,
+                type: m.type,
+                properties: m.properties
+            }, r.relationship_name, r.properties
+            """
+
+            result = await self.query(query, {"target_ids": target_ids})
+
+            if not result:
+                logger.info("No data returned for the supplied IDs")
+                return [], []
+
+            nodes_dict = {}
+            edges = []
+
+            for n_id, n_props, m_id, m_props, r_type, r_props_raw in result:
+                if n_props.get("properties"):
+                    try:
+                        additional_props = json.loads(n_props["properties"])
+                        n_props.update(additional_props)
+                        del n_props["properties"]
+                    except json.JSONDecodeError:
+                        logger.warning(f"Failed to parse properties JSON for node {n_id}")
+
+                if m_props.get("properties"):
+                    try:
+                        additional_props = json.loads(m_props["properties"])
+                        m_props.update(additional_props)
+                        del m_props["properties"]
+                    except json.JSONDecodeError:
+                        logger.warning(f"Failed to parse properties JSON for node {m_id}")
+
+                nodes_dict[n_id] = (n_id, n_props)
+                nodes_dict[m_id] = (m_id, m_props)
+
+                edge_props = {}
+                if r_props_raw:
+                    try:
+                        edge_props = json.loads(r_props_raw)
+                    except (json.JSONDecodeError, TypeError):
+                        logger.warning(f"Failed to parse edge properties for {n_id}->{m_id}")
+
+                source_id = edge_props.get("source_node_id", n_id)
+                target_id = edge_props.get("target_node_id", m_id)
+                edges.append((source_id, target_id, r_type, edge_props))
+
+            retrieval_time = time.time() - start_time
+            logger.info(
+                f"ID-filtered retrieval: {len(nodes_dict)} nodes and {len(edges)} edges in {retrieval_time:.2f}s"
+            )
+
+            return list(nodes_dict.values()), edges
+
+        except Exception as e:
+            logger.error(f"Error during ID-filtered graph data retrieval: {str(e)}")
+            raise
+
     async def get_graph_metrics(self, include_optional=False) -> Dict[str, Any]:
         """
         Get metrics on graph structure and connectivity.
@@ -1908,3 +2005,134 @@ class KuzuAdapter(GraphDBInterface):
         time_ids_list = [item[0] for item in time_nodes]
 
         return ", ".join(f"'{uid}'" for uid in time_ids_list)
+
+    async def get_triplets_batch(self, offset: int, limit: int) -> list[dict[str, Any]]:
+        """
+        Retrieve a batch of triplets (start_node, relationship, end_node) from the graph.
+
+        Parameters:
+        -----------
+            - offset (int): Number of triplets to skip before returning results.
+            - limit (int): Maximum number of triplets to return.
+
+        Returns:
+        --------
+            - list[dict[str, Any]]: A list of triplets, where each triplet is a dictionary
+              with keys: 'start_node', 'relationship_properties', 'end_node'.
+
+        Raises:
+        -------
+            - ValueError: If offset or limit are negative.
+            - Exception: Re-raises any exceptions from query execution.
+        """
+        if offset < 0:
+            raise ValueError(f"Offset must be non-negative, got {offset}")
+        if limit < 0:
+            raise ValueError(f"Limit must be non-negative, got {limit}")
+
+        query = """
+        MATCH (start_node:Node)-[relationship:EDGE]->(end_node:Node)
+        RETURN {
+            start_node: {
+                id: start_node.id,
+                name: start_node.name,
+                type: start_node.type,
+                properties: start_node.properties
+            },
+            relationship_properties: {
+                relationship_name: relationship.relationship_name,
+                properties: relationship.properties
+            },
+            end_node: {
+                id: end_node.id,
+                name: end_node.name,
+                type: end_node.type,
+                properties: end_node.properties
+            }
+        } AS triplet
+        SKIP $offset LIMIT $limit
+        """
+
+        try:
+            results = await self.query(query, {"offset": offset, "limit": limit})
+        except Exception as e:
+            logger.error(f"Failed to execute triplet query: {str(e)}")
+            logger.error(f"Query: {query}")
+            logger.error(f"Parameters: offset={offset}, limit={limit}")
+            raise
+
+        triplets = []
+        for idx, row in enumerate(results):
+            try:
+                if not row or len(row) == 0:
+                    logger.warning(f"Skipping empty row at index {idx} in triplet batch")
+                    continue
+
+                if not isinstance(row[0], dict):
+                    logger.warning(
+                        f"Skipping invalid row at index {idx}: expected dict, got {type(row[0])}"
+                    )
+                    continue
+
+                triplet = row[0]
+
+                if "start_node" not in triplet:
+                    logger.warning(f"Skipping triplet at index {idx}: missing 'start_node' key")
+                    continue
+
+                if not isinstance(triplet["start_node"], dict):
+                    logger.warning(f"Skipping triplet at index {idx}: 'start_node' is not a dict")
+                    continue
+
+                triplet["start_node"] = self._parse_node_properties(triplet["start_node"].copy())
+
+                if "relationship_properties" not in triplet:
+                    logger.warning(
+                        f"Skipping triplet at index {idx}: missing 'relationship_properties' key"
+                    )
+                    continue
+
+                if not isinstance(triplet["relationship_properties"], dict):
+                    logger.warning(
+                        f"Skipping triplet at index {idx}: 'relationship_properties' is not a dict"
+                    )
+                    continue
+
+                rel_props = triplet["relationship_properties"].copy()
+                relationship_name = rel_props.get("relationship_name") or ""
+
+                if rel_props.get("properties"):
+                    try:
+                        parsed_props = json.loads(rel_props["properties"])
+                        if isinstance(parsed_props, dict):
+                            rel_props.update(parsed_props)
+                            del rel_props["properties"]
+                        else:
+                            logger.warning(
+                                f"Parsed relationship properties is not a dict for triplet at index {idx}"
+                            )
+                    except (json.JSONDecodeError, TypeError) as e:
+                        logger.warning(
+                            f"Failed to parse relationship properties JSON for triplet at index {idx}: {e}"
+                        )
+
+                rel_props["relationship_name"] = relationship_name
+                triplet["relationship_properties"] = rel_props
+
+                if "end_node" not in triplet:
+                    logger.warning(f"Skipping triplet at index {idx}: missing 'end_node' key")
+                    continue
+
+                if not isinstance(triplet["end_node"], dict):
+                    logger.warning(f"Skipping triplet at index {idx}: 'end_node' is not a dict")
+                    continue
+
+                triplet["end_node"] = self._parse_node_properties(triplet["end_node"].copy())
+
+                triplets.append(triplet)
+
+            except Exception as e:
+                logger.error(f"Error processing triplet at index {idx}: {e}", exc_info=True)
+                continue
+
+        return triplets

cognee/infrastructure/databases/graph/neo4j_driver/Neo4jAuraDevDatasetDatabaseHandler.py

@@ -0,0 +1,168 @@
+import os
+import asyncio
+import requests
+import base64
+import hashlib
+from uuid import UUID
+from typing import Optional
+from cryptography.fernet import Fernet
+
+from cognee.infrastructure.databases.graph import get_graph_config
+from cognee.modules.users.models import User, DatasetDatabase
+from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface
+
+
+class Neo4jAuraDevDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
+    """
+    Handler for a quick development PoC integration of Cognee multi-user and permission mode with Neo4j Aura databases.
+    This handler creates a new Neo4j Aura instance for each Cognee dataset created.
+
+    Improvements needed to be production ready:
+    - Secret management for client credentials, currently secrets are encrypted and stored in the Cognee relational database,
+      a secret manager or a similar system should be used instead.
+
+    Quality of life improvements:
+    - Allow configuration of different Neo4j Aura plans and regions.
+    - Requests should be made async, currently a blocking requests library is used.
+    """
+
+    @classmethod
+    async def create_dataset(cls, dataset_id: Optional[UUID], user: Optional[User]) -> dict:
+        """
+        Create a new Neo4j Aura instance for the dataset. Return connection info that will be mapped to the dataset.
+
+        Args:
+            dataset_id: Dataset UUID
+            user: User object who owns the dataset and is making the request
+
+        Returns:
+            dict: Connection details for the created Neo4j instance
+
+        """
+        graph_config = get_graph_config()
+
+        if graph_config.graph_database_provider != "neo4j":
+            raise ValueError(
+                "Neo4jAuraDevDatasetDatabaseHandler can only be used with Neo4j graph database provider."
+            )
+
+        graph_db_name = f"{dataset_id}"
+
+        # Client credentials and encryption
+        client_id = os.environ.get("NEO4J_CLIENT_ID", None)
+        client_secret = os.environ.get("NEO4J_CLIENT_SECRET", None)
+        tenant_id = os.environ.get("NEO4J_TENANT_ID", None)
+        encryption_env_key = os.environ.get("NEO4J_ENCRYPTION_KEY", "test_key")
+        encryption_key = base64.urlsafe_b64encode(
+            hashlib.sha256(encryption_env_key.encode()).digest()
+        )
+        cipher = Fernet(encryption_key)
+
+        if client_id is None or client_secret is None or tenant_id is None:
+            raise ValueError(
+                "NEO4J_CLIENT_ID, NEO4J_CLIENT_SECRET, and NEO4J_TENANT_ID environment variables must be set to use Neo4j Aura DatasetDatabase Handling."
+            )
+
+        # Make the request with HTTP Basic Auth
+        def get_aura_token(client_id: str, client_secret: str) -> dict:
+            url = "https://api.neo4j.io/oauth/token"
+            data = {"grant_type": "client_credentials"}  # sent as application/x-www-form-urlencoded
+
+            resp = requests.post(url, data=data, auth=(client_id, client_secret))
+            resp.raise_for_status()  # raises if the request failed
+            return resp.json()
+
+        resp = get_aura_token(client_id, client_secret)
+
+        url = "https://api.neo4j.io/v1/instances"
+
+        headers = {
+            "accept": "application/json",
+            "Authorization": f"Bearer {resp['access_token']}",
+            "Content-Type": "application/json",
+        }
+
+        # TODO: Maybe we can allow **kwargs parameter forwarding for cases like these
+        #       Too allow different configurations between datasets
+        payload = {
+            "version": "5",
+            "region": "europe-west1",
+            "memory": "1GB",
+            "name": graph_db_name[
+                0:29
+            ],  # TODO: Find better name to name Neo4j instance within 30 character limit
+            "type": "professional-db",
+            "tenant_id": tenant_id,
+            "cloud_provider": "gcp",
+        }
+
+        response = requests.post(url, headers=headers, json=payload)
+
+        graph_db_name = "neo4j"  # Has to be 'neo4j' for Aura
+        graph_db_url = response.json()["data"]["connection_url"]
+        graph_db_key = resp["access_token"]
+        graph_db_username = response.json()["data"]["username"]
+        graph_db_password = response.json()["data"]["password"]
+
+        async def _wait_for_neo4j_instance_provisioning(instance_id: str, headers: dict):
+            # Poll until the instance is running
+            status_url = f"https://api.neo4j.io/v1/instances/{instance_id}"
+            status = ""
+            for attempt in range(30):  # Try for up to ~5 minutes
+                status_resp = requests.get(
+                    status_url, headers=headers
+                )  # TODO: Use async requests with httpx
+                status = status_resp.json()["data"]["status"]
+                if status.lower() == "running":
+                    return
+                await asyncio.sleep(10)
+            raise TimeoutError(
+                f"Neo4j instance '{graph_db_name}' did not become ready within 5 minutes. Status: {status}"
+            )
+
+        instance_id = response.json()["data"]["id"]
+        await _wait_for_neo4j_instance_provisioning(instance_id, headers)
+
+        encrypted_db_password_bytes = cipher.encrypt(graph_db_password.encode())
+        encrypted_db_password_string = encrypted_db_password_bytes.decode()
+
+        return {
+            "graph_database_name": graph_db_name,
+            "graph_database_url": graph_db_url,
+            "graph_database_provider": "neo4j",
+            "graph_database_key": graph_db_key,
+            "graph_dataset_database_handler": "neo4j_aura_dev",
+            "graph_database_connection_info": {
+                "graph_database_username": graph_db_username,
+                "graph_database_password": encrypted_db_password_string,
+            },
+        }
+
+    @classmethod
+    async def resolve_dataset_connection_info(
+        cls, dataset_database: DatasetDatabase
+    ) -> DatasetDatabase:
+        """
+        Resolve and decrypt connection info for the Neo4j dataset database.
+        In this case, decrypt the password stored in the database.
+
+        Args:
+            dataset_database: DatasetDatabase instance containing encrypted connection info.
+        """
+        encryption_env_key = os.environ.get("NEO4J_ENCRYPTION_KEY", "test_key")
+        encryption_key = base64.urlsafe_b64encode(
+            hashlib.sha256(encryption_env_key.encode()).digest()
+        )
+        cipher = Fernet(encryption_key)
+        graph_db_password = cipher.decrypt(
+            dataset_database.graph_database_connection_info["graph_database_password"].encode()
+        ).decode()
+
+        dataset_database.graph_database_connection_info["graph_database_password"] = (
+            graph_db_password
+        )
+        return dataset_database
+
+    @classmethod
+    async def delete_dataset(cls, dataset_database: DatasetDatabase):
+        pass