cognee 0.5.0.dev0__py3-none-any.whl → 0.5.1__py3-none-any.whl

cognee/eval_framework/Dockerfile

@@ -0,0 +1,29 @@
+FROM python:3.11-slim
+
+# Set environment variables
+ENV PIP_NO_CACHE_DIR=true
+ENV PATH="${PATH}:/root/.poetry/bin"
+ENV PYTHONPATH=/app
+ENV SKIP_MIGRATIONS=true
+
+# System dependencies
+RUN apt-get update && apt-get install -y \
+    gcc \
+    libpq-dev \
+    git \
+    curl \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+COPY pyproject.toml poetry.lock README.md /app/
+
+RUN pip install poetry
+
+RUN poetry config virtualenvs.create false
+
+RUN poetry install --extras distributed --extras evals --extras deepeval --no-root
+
+COPY cognee/ /app/cognee
+COPY distributed/ /app/distributed

cognee/eval_framework/answer_generation/answer_generation_executor.py

@@ -35,6 +35,16 @@ class AnswerGeneratorExecutor:
             retrieval_context = await retriever.get_context(query_text)
             search_results = await retriever.get_completion(query_text, retrieval_context)
 
+            ############
+            #:TODO This is a quick fix until we don't structure retriever results properly but lets not leave it like this...this is needed now due to the changed combined retriever structure..
+            if isinstance(retrieval_context, list):
+                retrieval_context = await retriever.convert_retrieved_objects_to_context(
+                    triplets=retrieval_context
+                )
+
+            if isinstance(search_results, str):
+                search_results = [search_results]
+            #############
             answer = {
                 "question": query_text,
                 "answer": search_results[0],

cognee/eval_framework/answer_generation/run_question_answering_module.py

@@ -35,7 +35,7 @@ async def create_and_insert_answers_table(questions_payload):
 
 
 async def run_question_answering(
-    params: dict, system_prompt="answer_simple_question.txt", top_k: Optional[int] = None
+    params: dict, system_prompt="answer_simple_question_benchmark.txt", top_k: Optional[int] = None
 ) -> List[dict]:
     if params.get("answering_questions"):
         logger.info("Question answering started...")

cognee/eval_framework/corpus_builder/task_getters/get_cascade_graph_tasks.py

@@ -8,7 +8,6 @@ from cognee.modules.users.models import User
 from cognee.shared.data_models import KnowledgeGraph
 from cognee.shared.utils import send_telemetry
 from cognee.tasks.documents import (
-    check_permissions_on_dataset,
     classify_documents,
     extract_chunks_from_documents,
 )
@@ -31,7 +30,6 @@ async def get_cascade_graph_tasks(
         cognee_config = get_cognify_config()
         default_tasks = [
             Task(classify_documents),
-            Task(check_permissions_on_dataset, user=user, permissions=["write"]),
             Task(
                 extract_chunks_from_documents, max_chunk_tokens=get_max_chunk_tokens()
             ),  # Extract text chunks based on the document type.

cognee/eval_framework/corpus_builder/task_getters/get_default_tasks_by_indices.py

@@ -30,8 +30,8 @@ async def get_no_summary_tasks(
     ontology_file_path=None,
 ) -> List[Task]:
     """Returns default tasks without summarization tasks."""
-    # Get base tasks (0=classify, 1=check_permissions, 2=extract_chunks)
-    base_tasks = await get_default_tasks_by_indices([0, 1, 2], chunk_size, chunker)
+    # Get base tasks (0=classify, 1=extract_chunks)
+    base_tasks = await get_default_tasks_by_indices([0, 1], chunk_size, chunker)
 
     ontology_adapter = RDFLibOntologyResolver(ontology_file=ontology_file_path)
 
@@ -51,8 +51,8 @@ async def get_just_chunks_tasks(
     chunk_size: int = None, chunker=TextChunker, user=None
 ) -> List[Task]:
     """Returns default tasks with only chunk extraction and data points addition."""
-    # Get base tasks (0=classify, 1=check_permissions, 2=extract_chunks)
-    base_tasks = await get_default_tasks_by_indices([0, 1, 2], chunk_size, chunker)
+    # Get base tasks (0=classify, 1=extract_chunks)
+    base_tasks = await get_default_tasks_by_indices([0, 1], chunk_size, chunker)
 
     add_data_points_task = Task(add_data_points, task_config={"batch_size": 10})
 

cognee/eval_framework/eval_config.py

@@ -14,7 +14,7 @@ class EvalConfig(BaseSettings):
 
     # Question answering params
     answering_questions: bool = True
-    qa_engine: str = "cognee_completion"  # Options: 'cognee_completion' or 'cognee_graph_completion' or 'cognee_graph_completion_cot' or 'cognee_graph_completion_context_extension'
+    qa_engine: str = "cognee_graph_completion"  # Options: 'cognee_completion' or 'cognee_graph_completion' or 'cognee_graph_completion_cot' or 'cognee_graph_completion_context_extension'
 
     # Evaluation params
     evaluating_answers: bool = True
@@ -25,7 +25,7 @@ class EvalConfig(BaseSettings):
         "EM",
         "f1",
     ]  # Use only 'correctness' for DirectLLM
-    deepeval_model: str = "gpt-5-mini"
+    deepeval_model: str = "gpt-4o-mini"
 
     # Metrics params
     calculate_metrics: bool = True

cognee/eval_framework/modal_run_eval.py

@@ -2,7 +2,6 @@ import modal
 import os
 import asyncio
 import datetime
-import hashlib
 import json
 from cognee.shared.logging_utils import get_logger
 from cognee.eval_framework.eval_config import EvalConfig
@@ -10,6 +9,9 @@ from cognee.eval_framework.corpus_builder.run_corpus_builder import run_corpus_b
 from cognee.eval_framework.answer_generation.run_question_answering_module import (
     run_question_answering,
 )
+import pathlib
+from os import path
+from modal import Image
 from cognee.eval_framework.evaluation.run_evaluation_module import run_evaluation
 from cognee.eval_framework.metrics_dashboard import create_dashboard
 
@@ -38,22 +40,19 @@ def read_and_combine_metrics(eval_params: dict) -> dict:
 
 app = modal.App("modal-run-eval")
 
-image = (
-    modal.Image.from_dockerfile(path="Dockerfile_modal", force_build=False)
-    .copy_local_file("pyproject.toml", "pyproject.toml")
-    .copy_local_file("poetry.lock", "poetry.lock")
-    .env(
-        {
-            "ENV": os.getenv("ENV"),
-            "LLM_API_KEY": os.getenv("LLM_API_KEY"),
-            "OPENAI_API_KEY": os.getenv("OPENAI_API_KEY"),
-        }
-    )
-    .pip_install("protobuf", "h2", "deepeval", "gdown", "plotly")
-)
+image = Image.from_dockerfile(
+    path=pathlib.Path(path.join(path.dirname(__file__), "Dockerfile")).resolve(),
+    force_build=False,
+).add_local_python_source("cognee")
 
 
-@app.function(image=image, concurrency_limit=10, timeout=86400, volumes={"/data": vol})
+@app.function(
+    image=image,
+    max_containers=10,
+    timeout=86400,
+    volumes={"/data": vol},
+    secrets=[modal.Secret.from_name("eval_secrets")],
+)
 async def modal_run_eval(eval_params=None):
     """Runs evaluation pipeline and returns combined metrics results."""
     if eval_params is None:
@@ -105,18 +104,7 @@ async def main():
     configs = [
         EvalConfig(
             task_getter_type="Default",
-            number_of_samples_in_corpus=10,
-            benchmark="HotPotQA",
-            qa_engine="cognee_graph_completion",
-            building_corpus_from_scratch=True,
-            answering_questions=True,
-            evaluating_answers=True,
-            calculate_metrics=True,
-            dashboard=True,
-        ),
-        EvalConfig(
-            task_getter_type="Default",
-            number_of_samples_in_corpus=10,
+            number_of_samples_in_corpus=25,
             benchmark="TwoWikiMultiHop",
             qa_engine="cognee_graph_completion",
             building_corpus_from_scratch=True,
@@ -127,7 +115,7 @@ async def main():
         ),
         EvalConfig(
             task_getter_type="Default",
-            number_of_samples_in_corpus=10,
+            number_of_samples_in_corpus=25,
             benchmark="Musique",
             qa_engine="cognee_graph_completion",
             building_corpus_from_scratch=True,

cognee/infrastructure/databases/dataset_database_handler/__init__.py

@@ -0,0 +1,3 @@
+from .dataset_database_handler_interface import DatasetDatabaseHandlerInterface
+from .supported_dataset_database_handlers import supported_dataset_database_handlers
+from .use_dataset_database_handler import use_dataset_database_handler

cognee/infrastructure/databases/dataset_database_handler/dataset_database_handler_interface.py

@@ -0,0 +1,80 @@
+from typing import Optional
+from uuid import UUID
+from abc import ABC, abstractmethod
+
+from cognee.modules.users.models.User import User
+from cognee.modules.users.models.DatasetDatabase import DatasetDatabase
+
+
+class DatasetDatabaseHandlerInterface(ABC):
+    @classmethod
+    @abstractmethod
+    async def create_dataset(cls, dataset_id: Optional[UUID], user: Optional[User]) -> dict:
+        """
+        Return a dictionary with database connection/resolution info for a graph or vector database for the given dataset.
+        Function can auto handle deploying of the actual database if needed, but is not necessary.
+        Only providing connection info is sufficient, this info will be mapped when trying to connect to the provided dataset in the future.
+        Needed for Cognee multi-tenant/multi-user and backend access control support.
+
+        Dictionary returned from this function will be used to create a DatasetDatabase row in the relational database.
+        From which internal mapping of dataset -> database connection info will be done.
+
+        The returned dictionary is stored verbatim in the relational database and is later passed to
+        resolve_dataset_connection_info() at connection time. For safe credential handling, prefer
+        returning only references to secrets or role identifiers, not plaintext credentials.
+
+        Each dataset needs to map to a unique graph or vector database when backend access control is enabled to facilitate a separation of concern for data.
+
+        Args:
+            dataset_id: UUID of the dataset if needed by the database creation logic
+            user: User object if needed by the database creation logic
+        Returns:
+            dict: Connection info for the created graph or vector database instance.
+        """
+        pass
+
+    @classmethod
+    async def resolve_dataset_connection_info(
+        cls, dataset_database: DatasetDatabase
+    ) -> DatasetDatabase:
+        """
+        Resolve runtime connection details for a dataset’s backing graph/vector database.
+        Function is intended to be overwritten to implement custom logic for resolving connection info.
+
+        This method is invoked right before the application opens a connection for a given dataset.
+        It receives the DatasetDatabase row that was persisted when create_dataset() ran and must
+        return a modified instance of DatasetDatabase with concrete connection parameters that the client/driver can use.
+        Do not update these new DatasetDatabase values in the relational database to avoid storing secure credentials.
+
+        In case of separate graph and vector database handlers, each handler should implement its own logic for resolving
+        connection info and only change parameters related to its appropriate database, the resolution function will then
+        be called one after another with the updated DatasetDatabase value from the previous function as the input.
+
+        Typical behavior:
+        - If the DatasetDatabase row already contains raw connection fields (e.g., host/port/db/user/password
+        or api_url/api_key), return them as-is.
+        - If the row stores only references (e.g., secret IDs, vault paths, cloud resource ARNs/IDs, IAM
+        roles, SSO tokens), resolve those references by calling the appropriate secret manager or provider
+        API to obtain short-lived credentials and assemble the final connection DatasetDatabase object.
+        - Do not persist any resolved or decrypted secrets back to the relational database. Return them only
+        to the caller.
+
+        Args:
+            dataset_database: DatasetDatabase row from the relational database
+        Returns:
+            DatasetDatabase: Updated instance with resolved connection info
+        """
+        return dataset_database
+
+    @classmethod
+    @abstractmethod
+    async def delete_dataset(cls, dataset_database: DatasetDatabase) -> None:
+        """
+        Delete the graph or vector database for the given dataset.
+        Function should auto handle deleting of the actual database or send a request to the proper service to delete/mark the database as not needed for the given dataset.
+        Needed for maintaining a database for Cognee multi-tenant/multi-user and backend access control.
+
+        Args:
+            dataset_database: DatasetDatabase row containing connection/resolution info for the graph or vector database to delete.
+        """
+        pass

cognee/infrastructure/databases/dataset_database_handler/supported_dataset_database_handlers.py

@@ -0,0 +1,18 @@
+from cognee.infrastructure.databases.graph.neo4j_driver.Neo4jAuraDevDatasetDatabaseHandler import (
+    Neo4jAuraDevDatasetDatabaseHandler,
+)
+from cognee.infrastructure.databases.vector.lancedb.LanceDBDatasetDatabaseHandler import (
+    LanceDBDatasetDatabaseHandler,
+)
+from cognee.infrastructure.databases.graph.kuzu.KuzuDatasetDatabaseHandler import (
+    KuzuDatasetDatabaseHandler,
+)
+
+supported_dataset_database_handlers = {
+    "neo4j_aura_dev": {
+        "handler_instance": Neo4jAuraDevDatasetDatabaseHandler,
+        "handler_provider": "neo4j",
+    },
+    "lancedb": {"handler_instance": LanceDBDatasetDatabaseHandler, "handler_provider": "lancedb"},
+    "kuzu": {"handler_instance": KuzuDatasetDatabaseHandler, "handler_provider": "kuzu"},
+}

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/graph/config.py

@@ -47,6 +47,7 @@ class GraphConfig(BaseSettings):
     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
@@ -97,6 +98,7 @@ class GraphConfig(BaseSettings):
             "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:
@@ -121,6 +123,7 @@ class GraphConfig(BaseSettings):
             "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

@@ -34,6 +34,7 @@ def create_graph_engine(
     graph_database_password="",
     graph_database_port="",
     graph_database_key="",
+    graph_dataset_database_handler="",
 ):
     """
     Create a graph engine based on the specified provider type.

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