cognee 0.2.1.dev3__tar.gz → 0.2.1.dev5__tar.gz

{cognee-0.2.1.dev3 → cognee-0.2.1.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 0.2.1.dev3
+Version: 0.2.1.dev5
 Summary: Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning.
 Project-URL: Homepage, https://www.cognee.ai
 Project-URL: Repository, https://github.com/topoteretes/cognee
@@ -115,8 +115,6 @@ Requires-Dist: langchain-text-splitters<1.0.0,>=0.3.2; extra == 'langchain'
 Requires-Dist: langsmith<1.0.0,>=0.2.3; extra == 'langchain'
 Provides-Extra: llama-index
 Requires-Dist: llama-index-core<0.13,>=0.12.11; extra == 'llama-index'
-Provides-Extra: milvus
-Requires-Dist: pymilvus<3,>=2.5.0; extra == 'milvus'
 Provides-Extra: mistral
 Requires-Dist: mistral-common<2,>=1.5.2; extra == 'mistral'
 Provides-Extra: neo4j

{cognee-0.2.1.dev3 → cognee-0.2.1.dev5}/cognee/api/v1/cognify/cognify.py

@@ -237,8 +237,12 @@ async def run_cognify_as_background_process(
     graph_db_config: dict = None,
     vector_db_config: dict = False,
 ):
+    # Convert dataset to list if it's a string
+    if isinstance(datasets, str):
+        datasets = [datasets]
+
     # Store pipeline status for all pipelines
-    pipeline_run_started_info = []
+    pipeline_run_started_info = {}
 
     async def handle_rest_of_the_run(pipeline_list):
         # Execute all provided pipelines one by one to avoid database write conflicts
@@ -269,7 +273,14 @@ async def run_cognify_as_background_process(
         )
 
         # Save dataset Pipeline run started info
-        pipeline_run_started_info.append(await anext(pipeline_run))
+        run_info = await anext(pipeline_run)
+        pipeline_run_started_info[run_info.dataset_id] = run_info
+
+        if pipeline_run_started_info[run_info.dataset_id].payload:
+            # Remove payload info to avoid serialization
+            # TODO: Handle payload serialization
+            pipeline_run_started_info[run_info.dataset_id].payload = []
+
         pipeline_list.append(pipeline_run)
 
     # Send all started pipelines to execute one by one in background

{cognee-0.2.1.dev3 → cognee-0.2.1.dev5}/cognee/api/v1/cognify/routers/get_cognify_router.py

@@ -32,7 +32,6 @@ logger = get_logger("api.cognify")
 class CognifyPayloadDTO(InDTO):
     datasets: Optional[List[str]] = None
     dataset_ids: Optional[List[UUID]] = None
-    graph_model: Optional[BaseModel] = KnowledgeGraph
     run_in_background: Optional[bool] = False
 
 
@@ -96,7 +95,7 @@ def get_cognify_router() -> APIRouter:
             datasets = payload.dataset_ids if payload.dataset_ids else payload.datasets
 
             cognify_run = await cognee_cognify(
-                datasets, user, payload.graph_model, run_in_background=payload.run_in_background
+                datasets, user, run_in_background=payload.run_in_background
             )
 
             return cognify_run

{cognee-0.2.1.dev3 → cognee-0.2.1.dev5}/cognee/api/v1/config/config.py

@@ -15,7 +15,7 @@ class config:
     @staticmethod
     def system_root_directory(system_root_directory: str):
         base_config = get_base_config()
-        base_config.system_root_directory = os.path.join(system_root_directory, ".cognee_system")
+        base_config.system_root_directory = system_root_directory
 
         databases_directory_path = os.path.join(base_config.system_root_directory, "databases")
 
@@ -24,13 +24,7 @@ class config:
 
         graph_config = get_graph_config()
         graph_file_name = graph_config.graph_filename
-        # For Kuzu v0.11.0+, use single-file database with .kuzu extension
-        if graph_config.graph_database_provider.lower() == "kuzu":
-            graph_config.graph_file_path = os.path.join(
-                databases_directory_path, f"{graph_file_name}.kuzu"
-            )
-        else:
-            graph_config.graph_file_path = os.path.join(databases_directory_path, graph_file_name)
+        graph_config.graph_file_path = os.path.join(databases_directory_path, graph_file_name)
 
         vector_config = get_vectordb_config()
         if vector_config.vector_db_provider == "lancedb":
@@ -39,7 +33,7 @@ class config:
     @staticmethod
     def data_root_directory(data_root_directory: str):
         base_config = get_base_config()
-        base_config.data_root_directory = os.path.join(data_root_directory, ".data_storage")
+        base_config.data_root_directory = data_root_directory
 
     @staticmethod
     def monitoring_tool(monitoring_tool: object):

{cognee-0.2.1.dev3 → cognee-0.2.1.dev5}/cognee/api/v1/datasets/routers/get_datasets_router.py

@@ -331,11 +331,19 @@ def get_datasets_router() -> APIRouter:
         ## Error Codes
         - **500 Internal Server Error**: Error retrieving status information
         """
-        from cognee.modules.data.methods import get_dataset_status
+        from cognee.api.v1.datasets.datasets import datasets as cognee_datasets
 
-        dataset_status = await get_dataset_status(datasets, user.id)
+        try:
+            # Verify user has permission to read dataset
+            authorized_datasets = await get_authorized_existing_datasets(datasets, "read", user)
+
+            datasets_statuses = await cognee_datasets.get_status(
+                [dataset.id for dataset in authorized_datasets]
+            )
 
-        return dataset_status
+            return datasets_statuses
+        except Exception as error:
+            return JSONResponse(status_code=409, content={"error": str(error)})
 
     @router.get("/{dataset_id}/data/{data_id}/raw", response_class=FileResponse)
     async def get_raw_data(

{cognee-0.2.1.dev3 → cognee-0.2.1.dev5}/cognee/api/v1/delete/delete.py

@@ -1,100 +1,91 @@
-import os
-import hashlib
 from uuid import UUID
-from io import BytesIO
 from sqlalchemy import select
-from typing import Union, BinaryIO, List
 from sqlalchemy.sql import delete as sql_delete
 
 from cognee.infrastructure.engine import DataPoint
-from cognee.infrastructure.files.storage import get_file_storage
 from cognee.infrastructure.databases.graph import get_graph_engine
+
+from cognee.modules.users.models import User
+
 from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.infrastructure.databases.relational import get_relational_engine
-from cognee.modules.ingestion import classify
-from cognee.modules.users.models import User
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.data.models import Data, DatasetData, Dataset
 from cognee.modules.graph.utils.convert_node_to_data_point import get_all_subclasses
 from cognee.modules.users.methods import get_default_user
 from cognee.modules.data.methods import get_authorized_existing_datasets
 from cognee.context_global_variables import set_database_global_context_variables
+
 from .exceptions import DocumentNotFoundError, DatasetNotFoundError, DocumentSubgraphNotFoundError
 
 logger = get_logger()
 
 
-def get_text_content_hash(text: str) -> str:
-    encoded_text = text.encode("utf-8")
-    return hashlib.md5(encoded_text).hexdigest()
-
-
 async def delete(
-    data: Union[BinaryIO, List[BinaryIO], str, List[str]],
-    dataset_name: str = "main_dataset",
-    dataset_id: UUID = None,
+    data_id: UUID,
+    dataset_id: UUID,
     mode: str = "soft",
     user: User = None,
 ):
-    """Delete a document and all its related nodes from both relational and graph databases.
+    """Delete data by its ID from the specified dataset.
 
     Args:
-        data: The data to delete (file, URL, or text)
-        dataset_name: Name of the dataset to delete from
+        data_id: The UUID of the data to delete
+        dataset_id: The UUID of the dataset containing the data
         mode: "soft" (default) or "hard" - hard mode also deletes degree-one entity nodes
         user: User doing the operation, if none default user will be used.
-    """
 
+    Returns:
+        Dict with deletion results
+
+    Raises:
+        DocumentNotFoundError: If data is not found
+        DatasetNotFoundError: If dataset is not found
+        PermissionDeniedError: If user doesn't have delete permission on dataset
+    """
     if user is None:
         user = await get_default_user()
 
-    # Verify user has permission to work with given dataset. If dataset_id is given use it, if not use dataset_name
-    dataset = await get_authorized_existing_datasets(
-        [dataset_id] if dataset_id else [dataset_name], "delete", user
-    )
+    # Verify user has delete permission on the dataset
+    dataset_list = await get_authorized_existing_datasets([dataset_id], "delete", user)
+
+    if not dataset_list:
+        raise DatasetNotFoundError(f"Dataset not found or access denied: {dataset_id}")
+
+    dataset = dataset_list[0]
 
     # Will only be used if ENABLE_BACKEND_ACCESS_CONTROL is set to True
-    await set_database_global_context_variables(dataset[0].id, dataset[0].owner_id)
-
-    # Handle different input types
-    if isinstance(data, str):
-        if data.startswith("file://") or data.startswith("/"):  # It's a file path
-            full_file_path = data.replace("file://", "")
-
-            file_dir = os.path.dirname(full_file_path)
-            file_path = os.path.basename(full_file_path)
-
-            file_storage = get_file_storage(file_dir)
-
-            async with file_storage.open(file_path, mode="rb") as file:
-                classified_data = classify(file)
-                content_hash = classified_data.get_metadata()["content_hash"]
-                return await delete_single_document(content_hash, dataset[0].id, mode)
-        elif data.startswith("http"):  # It's a URL
-            import requests
-
-            response = requests.get(data)
-            response.raise_for_status()
-            file_data = BytesIO(response.content)
-            classified_data = classify(file_data)
-            content_hash = classified_data.get_metadata()["content_hash"]
-            return await delete_single_document(content_hash, dataset[0].id, mode)
-        else:  # It's a text string
-            content_hash = get_text_content_hash(data)
-            classified_data = classify(data)
-            return await delete_single_document(content_hash, dataset[0].id, mode)
-    elif isinstance(data, list):
-        # Handle list of inputs sequentially
-        results = []
-        for item in data:
-            result = await delete(item, dataset_name, dataset[0].id, mode, user=user)
-            results.append(result)
-        return {"status": "success", "message": "Multiple documents deleted", "results": results}
-    else:  # It's already a BinaryIO
-        data.seek(0)  # Ensure we're at the start of the file
-        classified_data = classify(data)
-        content_hash = classified_data.get_metadata()["content_hash"]
-        return await delete_single_document(content_hash, dataset[0].id, mode)
+
+    await set_database_global_context_variables(dataset.id, dataset.owner_id)
+
+    # Get the data record and verify it exists and belongs to the dataset
+    db_engine = get_relational_engine()
+    async with db_engine.get_async_session() as session:
+        # Check if data exists
+        data_point = (
+            await session.execute(select(Data).filter(Data.id == data_id))
+        ).scalar_one_or_none()
+
+        if data_point is None:
+            raise DocumentNotFoundError(f"Data not found with ID: {data_id}")
+
+        # Check if data belongs to the specified dataset
+        dataset_data_link = (
+            await session.execute(
+                select(DatasetData).filter(
+                    DatasetData.data_id == data_id, DatasetData.dataset_id == dataset_id
+                )
+            )
+        ).scalar_one_or_none()
+
+        if dataset_data_link is None:
+            raise DocumentNotFoundError(f"Data {data_id} not found in dataset {dataset_id}")
+
+        # Get the content hash for deletion
+        content_hash = data_point.content_hash
+
+    # Use the existing comprehensive deletion logic
+    return await delete_single_document(content_hash, dataset.id, mode)
 
 
 async def delete_single_document(content_hash: str, dataset_id: UUID = None, mode: str = "soft"):

cognee-0.2.1.dev5/cognee/api/v1/delete/routers/get_delete_router.py

@@ -0,0 +1,49 @@
+from fastapi import Depends
+from fastapi.responses import JSONResponse
+from fastapi import APIRouter
+from uuid import UUID
+from cognee.shared.logging_utils import get_logger
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_authenticated_user
+
+logger = get_logger()
+
+
+def get_delete_router() -> APIRouter:
+    router = APIRouter()
+
+    @router.delete("", response_model=None)
+    async def delete(
+        data_id: UUID,
+        dataset_id: UUID,
+        mode: str = "soft",
+        user: User = Depends(get_authenticated_user),
+    ):
+        """Delete data by its ID from the specified dataset.
+
+        Args:
+            data_id: The UUID of the data to delete
+            dataset_id: The UUID of the dataset containing the data
+            mode: "soft" (default) or "hard" - hard mode also deletes degree-one entity nodes
+            user: Authenticated user
+
+        Returns:
+            JSON response indicating success or failure
+
+        """
+        from cognee.api.v1.delete import delete as cognee_delete
+
+        try:
+            result = await cognee_delete(
+                data_id=data_id,
+                dataset_id=dataset_id,
+                mode=mode,
+                user=user,
+            )
+            return result
+
+        except Exception as error:
+            logger.error(f"Error during deletion by data_id: {str(error)}")
+            return JSONResponse(status_code=409, content={"error": str(error)})
+
+    return router

{cognee-0.2.1.dev3 → cognee-0.2.1.dev5}/cognee/context_global_variables.py

@@ -5,8 +5,8 @@ from uuid import UUID
 
 from cognee.base_config import get_base_config
 from cognee.infrastructure.databases.utils import get_or_create_dataset_database
-from cognee.modules.users.methods import get_user
 from cognee.infrastructure.files.storage.config import file_storage_config
+from cognee.modules.users.methods import get_user
 
 # Note: ContextVar allows us to use different graph db configurations in Cognee
 #       for different async tasks, threads and processes
@@ -46,13 +46,15 @@ async def set_database_global_context_variables(dataset: Union[str, UUID], user_
     data_root_directory = os.path.join(
         base_config.data_root_directory, str(user.tenant_id or user.id)
     )
-    system_directory_path = os.path.join(
+    databases_directory_path = os.path.join(
         base_config.system_root_directory, "databases", str(user.id)
     )
 
     # Set vector and graph database configuration based on dataset database information
     vector_config = {
-        "vector_db_url": os.path.join(system_directory_path, dataset_database.vector_database_name),
+        "vector_db_url": os.path.join(
+            databases_directory_path, dataset_database.vector_database_name
+        ),
         "vector_db_key": "",
         "vector_db_provider": "lancedb",
     }
@@ -60,7 +62,7 @@ async def set_database_global_context_variables(dataset: Union[str, UUID], user_
     graph_config = {
         "graph_database_provider": "kuzu",
         "graph_file_path": os.path.join(
-            system_directory_path, f"{dataset_database.graph_database_name}.kuzu"
+            databases_directory_path, dataset_database.graph_database_name
         ),
     }
 

{cognee-0.2.1.dev3 → cognee-0.2.1.dev5}/cognee/infrastructure/databases/graph/config.py

@@ -50,19 +50,18 @@ class GraphConfig(BaseSettings):
     @pydantic.model_validator(mode="after")
     def fill_derived(cls, values):
         provider = values.graph_database_provider.lower()
-        # Set filename based on graph database provider if no filename is provided
+
+        # Set default filename if no filename is provided
         if not values.graph_filename:
             values.graph_filename = f"cognee_graph_{provider}"
+
         # Set file path based on graph database provider if no file path is provided
         if not values.graph_file_path:
             base_config = get_base_config()
 
-            base = os.path.join(base_config.system_root_directory, "databases")
-            # For Kuzu v0.11.0+, use single-file database with .kuzu extension
-            if provider == "kuzu":
-                values.graph_file_path = os.path.join(base, f"{values.graph_filename}.kuzu")
-            else:
-                values.graph_file_path = os.path.join(base, values.graph_filename)
+            databases_directory_path = os.path.join(base_config.system_root_directory, "databases")
+            values.graph_file_path = os.path.join(databases_directory_path, values.graph_filename)
+
         return values
 
     def to_dict(self) -> dict:

{cognee-0.2.1.dev3 → cognee-0.2.1.dev5}/cognee/infrastructure/databases/graph/kuzu/adapter.py

@@ -54,8 +54,8 @@ class KuzuAdapter(GraphDBInterface):
 
                 self.db = Database(
                     self.temp_graph_file,
-                    buffer_pool_size=256 * 1024 * 1024,  # 256MB buffer pool
-                    max_db_size=1024 * 1024 * 1024,
+                    buffer_pool_size=2048 * 1024 * 1024,  # 2048MB buffer pool
+                    max_db_size=4096 * 1024 * 1024,
                 )
             else:
                 # Ensure the parent directory exists before creating the database
@@ -1452,9 +1452,13 @@ class KuzuAdapter(GraphDBInterface):
             db_name = os.path.basename(self.db_path)
             file_storage = get_file_storage(db_dir)
 
-            if await file_storage.file_exists(db_name):
-                await file_storage.remove_all()
-                logger.info(f"Deleted Kuzu database files at {self.db_path}")
+            if await file_storage.is_file(db_name):
+                await file_storage.remove(db_name)
+                await file_storage.remove(f"{db_name}.lock")
+            else:
+                await file_storage.remove_all(db_name)
+
+            logger.info(f"Deleted Kuzu database files at {self.db_path}")
 
         except Exception as e:
             logger.error(f"Failed to delete graph data: {e}")

{cognee-0.2.1.dev3 → cognee-0.2.1.dev5}/cognee/infrastructure/databases/graph/memgraph/memgraph_adapter.py

@@ -381,12 +381,19 @@ class MemgraphAdapter(GraphDBInterface):
 
             The result of the edge addition operation, including relationship details.
         """
+
+        exists = await asyncio.gather(self.has_node(str(from_node)), self.has_node(str(to_node)))
+
+        if not all(exists):
+            return None
+
         serialized_properties = self.serialize_properties(edge_properties or {})
 
         query = dedent(
             f"""\
             MATCH (from_node {{id: $from_node}}),
                   (to_node {{id: $to_node}})
+            WHERE from_node IS NOT NULL AND to_node IS NOT NULL
             MERGE (from_node)-[r:{relationship_name}]->(to_node)
             ON CREATE SET r += $properties, r.updated_at = timestamp()
             ON MATCH SET r += $properties, r.updated_at = timestamp()

{cognee-0.2.1.dev3 → cognee-0.2.1.dev5}/cognee/infrastructure/databases/graph/networkx/adapter.py

@@ -679,11 +679,11 @@ class NetworkXAdapter(GraphDBInterface):
             )  # Assuming self.filename is defined elsewhere and holds the default graph file path
         try:
             file_dir_path = os.path.dirname(file_path)
-            file_path = os.path.basename(file_path)
+            file_name = os.path.basename(file_path)
 
             file_storage = get_file_storage(file_dir_path)
 
-            await file_storage.remove(file_path)
+            await file_storage.remove(file_name)
 
             self.graph = None
             logger.info("Graph deleted successfully.")