cognee 0.2.3.dev1__tar.gz → 0.2.4__tar.gz

{cognee-0.2.3.dev1 → cognee-0.2.4}/.env.template

@@ -6,7 +6,7 @@
 #       Default graph database    : Kuzu
 #
 #       These default databases are all file-based, so no extra setup is needed
-#       for local use.
+#       for local use. The data by default will be stored in your .venv
 ###############################################################################
 
 ################################################################################
@@ -16,7 +16,7 @@
 STRUCTURED_OUTPUT_FRAMEWORK="instructor"
 
 LLM_API_KEY="your_api_key"
-LLM_MODEL="openai/gpt-4o-mini"
+LLM_MODEL="openai/gpt-5-mini"
 LLM_PROVIDER="openai"
 LLM_ENDPOINT=""
 LLM_API_VERSION=""
@@ -33,11 +33,20 @@ EMBEDDING_MAX_TOKENS=8191
 
 # If using BAML structured output these env variables will be used
 BAML_LLM_PROVIDER=openai
-BAML_LLM_MODEL="gpt-4o-mini"
+BAML_LLM_MODEL="gpt-5-mini"
 BAML_LLM_ENDPOINT=""
 BAML_LLM_API_KEY="your_api_key"
 BAML_LLM_API_VERSION=""
 
+
+################################################################################
+#  📂 ROOT DIRECTORY FOR DATABASES
+################################################################################
+# Set up the Cognee system directory. Cognee will store system files and databases here.
+# Useful for setting root directory inside docker and also to avoid storing the databases in .venv
+# DATA_ROOT_DIRECTORY='/Users/<user>/Desktop/cognee/.cognee_data/'
+# SYSTEM_ROOT_DIRECTORY='/Users/<user>/Desktop/cognee/.cognee_system/'
+
 ################################################################################
 # 🗄️ Relational database settings
 ################################################################################
@@ -85,12 +94,7 @@ VECTOR_DB_PROVIDER="lancedb"
 VECTOR_DB_URL=
 VECTOR_DB_KEY=
 
-################################################################################
-#  📂 ROOT DIRECTORY IF USING COGNEE LIB INSIDE A DOCKER
-################################################################################
-# Set up the Cognee system directory. Cognee will store system files and databases here.
-# DATA_ROOT_DIRECTORY='/Users/<user>/Desktop/cognee/.cognee_data/'
-# SYSTEM_ROOT_DIRECTORY='/Users/<user>/Desktop/cognee/.cognee_system/'
+
 
 ################################################################################
 #  🔄  MIGRATION (RELATIONAL → GRAPH) SETTINGS
@@ -157,8 +161,8 @@ LITELLM_LOG="ERROR"
 # Uncomment + fill values to switch.
 
 ########## Azure OpenAI #######################################################
-#LLM_MODEL="azure/gpt-4o-mini"
-#LLM_ENDPOINT="https://DNS.azure.com/openai/deployments/gpt-4o-mini"
+#LLM_MODEL="azure/gpt-5-mini"
+#LLM_ENDPOINT="https://DNS.azure.com/openai/deployments/gpt-5-mini"
 #LLM_API_KEY="<<TALK TO YOUR AZURE GUY"
 #LLM_API_VERSION="2024-12-01-preview"
 

{cognee-0.2.3.dev1 → cognee-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 0.2.3.dev1
+Version: 0.2.4
 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
@@ -31,13 +31,13 @@ Requires-Dist: kuzu==0.11.0
 Requires-Dist: lancedb<1.0.0,>=0.24.0
 Requires-Dist: langfuse<3,>=2.32.0
 Requires-Dist: limits<5,>=4.4.1
-Requires-Dist: litellm<1.71.0,>=1.57.4
+Requires-Dist: litellm<2.0.0,>=1.71.0
 Requires-Dist: matplotlib<4,>=3.8.3
 Requires-Dist: networkx<4,>=3.4.2
 Requires-Dist: nltk<4.0.0,>=3.9.1
 Requires-Dist: numpy<=4.0.0,>=1.26.4
 Requires-Dist: onnxruntime<2.0.0,>=1.0.0
-Requires-Dist: openai<2,>=1.80.1
+Requires-Dist: openai<1.99.9,>=1.80.1
 Requires-Dist: pandas<3.0.0,>=2.2.2
 Requires-Dist: pre-commit<5,>=4.0.1
 Requires-Dist: pydantic-settings<3,>=2.2.1
@@ -46,6 +46,7 @@ Requires-Dist: pylance<1.0.0,>=0.22.0
 Requires-Dist: pympler<2.0.0,>=1.1
 Requires-Dist: pypdf<6.0.0,>=4.1.0
 Requires-Dist: python-dotenv<2.0.0,>=1.0.1
+Requires-Dist: python-magic-bin<0.5; platform_system == 'Windows'
 Requires-Dist: python-multipart<1.0.0,>=0.0.20
 Requires-Dist: rdflib<7.2.0,>=7.1.4
 Requires-Dist: s3fs[boto3]==2025.3.2
@@ -74,7 +75,7 @@ Requires-Dist: tree-sitter<0.25,>=0.24.0; extra == 'codegraph'
 Provides-Extra: debug
 Requires-Dist: debugpy<2.0.0,>=1.8.9; extra == 'debug'
 Provides-Extra: deepeval
-Requires-Dist: deepeval<3,>=2.0.1; extra == 'deepeval'
+Requires-Dist: deepeval<4,>=3.0.1; extra == 'deepeval'
 Provides-Extra: dev
 Requires-Dist: coverage<8,>=7.3.2; extra == 'dev'
 Requires-Dist: deptry<0.21,>=0.20.0; extra == 'dev'
@@ -230,7 +231,8 @@ Your contributions are at the core of making this a true open source project. An
 
 ## 📦 Installation
 
-You can install Cognee using either **uv**, **pip**, **poetry** or any other python package manager.
+You can install Cognee using either **pip**, **poetry**, **uv** or any other python package manager.
+
 Cognee supports Python 3.10 to 3.13
 
 ### With pip

{cognee-0.2.3.dev1 → cognee-0.2.4}/README.md

@@ -91,7 +91,8 @@ Your contributions are at the core of making this a true open source project. An
 
 ## 📦 Installation
 
-You can install Cognee using either **uv**, **pip**, **poetry** or any other python package manager.
+You can install Cognee using either **pip**, **poetry**, **uv** or any other python package manager.
+
 Cognee supports Python 3.10 to 3.13
 
 ### With pip

cognee-0.2.4/cognee/__main__.py

@@ -0,0 +1,4 @@
+from cognee.cli._cognee import main
+
+if __name__ == "__main__":
+    main()

{cognee-0.2.3.dev1 → cognee-0.2.4}/cognee/api/v1/add/add.py

@@ -1,9 +1,15 @@
 from uuid import UUID
 from typing import Union, BinaryIO, List, Optional
 
-from cognee.modules.pipelines import Task
 from cognee.modules.users.models import User
-from cognee.modules.pipelines import cognee_pipeline
+from cognee.modules.pipelines import Task, run_pipeline
+from cognee.modules.pipelines.layers.resolve_authorized_user_dataset import (
+    resolve_authorized_user_dataset,
+)
+from cognee.modules.pipelines.layers.reset_dataset_pipeline_run_status import (
+    reset_dataset_pipeline_run_status,
+)
+from cognee.modules.engine.operations.setup import setup
 from cognee.tasks.ingestion import ingest_data, resolve_data_directories
 
 
@@ -128,11 +134,11 @@ async def add(
 
         Optional:
         - LLM_PROVIDER: "openai" (default), "anthropic", "gemini", "ollama"
-        - LLM_MODEL: Model name (default: "gpt-4o-mini")
+        - LLM_MODEL: Model name (default: "gpt-5-mini")
         - DEFAULT_USER_EMAIL: Custom default user email
         - DEFAULT_USER_PASSWORD: Custom default user password
         - VECTOR_DB_PROVIDER: "lancedb" (default), "chromadb", "pgvector"
-        - GRAPH_DATABASE_PROVIDER: "kuzu" (default), "neo4j", "networkx"
+        - GRAPH_DATABASE_PROVIDER: "kuzu" (default), "neo4j"
 
     """
     tasks = [
@@ -140,11 +146,17 @@ async def add(
         Task(ingest_data, dataset_name, user, node_set, dataset_id, preferred_loaders),
     ]
 
+    await setup()
+
+    user, authorized_dataset = await resolve_authorized_user_dataset(dataset_id, dataset_name, user)
+
+    await reset_dataset_pipeline_run_status(authorized_dataset.id, user)
+
     pipeline_run_info = None
 
-    async for run_info in cognee_pipeline(
+    async for run_info in run_pipeline(
         tasks=tasks,
-        datasets=dataset_id if dataset_id else dataset_name,
+        datasets=[authorized_dataset.id],
         data=data,
         user=user,
         pipeline_name="add_pipeline",

{cognee-0.2.3.dev1 → cognee-0.2.4}/cognee/api/v1/cognify/code_graph_pipeline.py

@@ -40,8 +40,14 @@ async def run_code_graph_pipeline(repo_path, include_docs=False):
     user = await get_default_user()
     detailed_extraction = True
 
+    # Multi-language support: allow passing supported_languages
+    supported_languages = None  # defer to task defaults
     tasks = [
-        Task(get_repo_file_dependencies, detailed_extraction=detailed_extraction),
+        Task(
+            get_repo_file_dependencies,
+            detailed_extraction=detailed_extraction,
+            supported_languages=supported_languages,
+        ),
         # Task(summarize_code, task_config={"batch_size": 500}), # This task takes a long time to complete
         Task(add_data_points, task_config={"batch_size": 30}),
     ]

{cognee-0.2.3.dev1 → cognee-0.2.4}/cognee/api/v1/cognify/cognify.py

@@ -7,12 +7,10 @@ from cognee.shared.logging_utils import get_logger
 from cognee.shared.data_models import KnowledgeGraph
 from cognee.infrastructure.llm import get_max_chunk_tokens
 
-from cognee.modules.pipelines import cognee_pipeline
+from cognee.modules.pipelines import run_pipeline
 from cognee.modules.pipelines.tasks.task import Task
 from cognee.modules.chunking.TextChunker import TextChunker
 from cognee.modules.ontology.rdf_xml.OntologyResolver import OntologyResolver
-from cognee.modules.pipelines.models.PipelineRunInfo import PipelineRunCompleted, PipelineRunErrored
-from cognee.modules.pipelines.queues.pipeline_run_info_queues import push_to_queue
 from cognee.modules.users.models import User
 
 from cognee.tasks.documents import (
@@ -23,6 +21,7 @@ from cognee.tasks.documents import (
 from cognee.tasks.graph import extract_graph_from_data
 from cognee.tasks.storage import add_data_points
 from cognee.tasks.summarization import summarize_text
+from cognee.modules.pipelines.layers.pipeline_execution_mode import get_pipeline_executor
 
 logger = get_logger("cognify")
 
@@ -40,6 +39,7 @@ async def cognify(
     graph_db_config: dict = None,
     run_in_background: bool = False,
     incremental_loading: bool = True,
+    custom_prompt: Optional[str] = None,
 ):
     """
     Transform ingested data into a structured knowledge graph.
@@ -91,7 +91,7 @@ async def cognify(
                 - LangchainChunker: Recursive character splitting with overlap
                 Determines how documents are segmented for processing.
         chunk_size: Maximum tokens per chunk. Auto-calculated based on LLM if None.
-                   Formula: min(embedding_max_tokens, llm_max_tokens // 2)
+                   Formula: min(embedding_max_completion_tokens, llm_max_completion_tokens // 2)
                    Default limits: ~512-8192 tokens depending on models.
                    Smaller chunks = more granular but potentially fragmented knowledge.
         ontology_file_path: Path to RDF/OWL ontology file for domain-specific entity types.
@@ -102,6 +102,10 @@ async def cognify(
                           If False, waits for completion before returning.
                           Background mode recommended for large datasets (>100MB).
                           Use pipeline_run_id from return value to monitor progress.
+        custom_prompt: Optional custom prompt string to use for entity extraction and graph generation.
+                      If provided, this prompt will be used instead of the default prompts for
+                      knowledge graph extraction. The prompt should guide the LLM on how to
+                      extract entities and relationships from the text content.
 
     Returns:
         Union[dict, list[PipelineRunInfo]]:
@@ -178,115 +182,24 @@ async def cognify(
         - LLM_RATE_LIMIT_ENABLED: Enable rate limiting (default: False)
         - LLM_RATE_LIMIT_REQUESTS: Max requests per interval (default: 60)
     """
-    tasks = await get_default_tasks(user, graph_model, chunker, chunk_size, ontology_file_path)
-
-    if run_in_background:
-        return await run_cognify_as_background_process(
-            tasks=tasks,
-            user=user,
-            datasets=datasets,
-            vector_db_config=vector_db_config,
-            graph_db_config=graph_db_config,
-            incremental_loading=incremental_loading,
-        )
-    else:
-        return await run_cognify_blocking(
-            tasks=tasks,
-            user=user,
-            datasets=datasets,
-            vector_db_config=vector_db_config,
-            graph_db_config=graph_db_config,
-            incremental_loading=incremental_loading,
-        )
-
+    tasks = await get_default_tasks(
+        user, graph_model, chunker, chunk_size, ontology_file_path, custom_prompt
+    )
 
-async def run_cognify_blocking(
-    tasks,
-    user,
-    datasets,
-    graph_db_config: dict = None,
-    vector_db_config: dict = False,
-    incremental_loading: bool = True,
-):
-    total_run_info = {}
+    # By calling get pipeline executor we get a function that will have the run_pipeline run in the background or a function that we will need to wait for
+    pipeline_executor_func = get_pipeline_executor(run_in_background=run_in_background)
 
-    async for run_info in cognee_pipeline(
+    # Run the run_pipeline in the background or blocking based on executor
+    return await pipeline_executor_func(
+        pipeline=run_pipeline,
         tasks=tasks,
-        datasets=datasets,
         user=user,
-        pipeline_name="cognify_pipeline",
-        graph_db_config=graph_db_config,
+        datasets=datasets,
         vector_db_config=vector_db_config,
+        graph_db_config=graph_db_config,
         incremental_loading=incremental_loading,
-    ):
-        if run_info.dataset_id:
-            total_run_info[run_info.dataset_id] = run_info
-        else:
-            total_run_info = run_info
-
-    return total_run_info
-
-
-async def run_cognify_as_background_process(
-    tasks,
-    user,
-    datasets,
-    graph_db_config: dict = None,
-    vector_db_config: dict = False,
-    incremental_loading: bool = True,
-):
-    # 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 = {}
-
-    async def handle_rest_of_the_run(pipeline_list):
-        # Execute all provided pipelines one by one to avoid database write conflicts
-        # TODO: Convert to async gather task instead of for loop when Queue mechanism for database is created
-        for pipeline in pipeline_list:
-            while True:
-                try:
-                    pipeline_run_info = await anext(pipeline)
-
-                    push_to_queue(pipeline_run_info.pipeline_run_id, pipeline_run_info)
-
-                    if isinstance(pipeline_run_info, PipelineRunCompleted) or isinstance(
-                        pipeline_run_info, PipelineRunErrored
-                    ):
-                        break
-                except StopAsyncIteration:
-                    break
-
-    # Start all pipelines to get started status
-    pipeline_list = []
-    for dataset in datasets:
-        pipeline_run = cognee_pipeline(
-            tasks=tasks,
-            user=user,
-            datasets=dataset,
-            pipeline_name="cognify_pipeline",
-            graph_db_config=graph_db_config,
-            vector_db_config=vector_db_config,
-            incremental_loading=incremental_loading,
-        )
-
-        # Save dataset Pipeline run started info
-        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
-    asyncio.create_task(handle_rest_of_the_run(pipeline_list=pipeline_list))
-
-    return pipeline_run_started_info
+        pipeline_name="cognify_pipeline",
+    )
 
 
 async def get_default_tasks(  # TODO: Find out a better way to do this (Boris's comment)
@@ -295,6 +208,7 @@ async def get_default_tasks(  # TODO: Find out a better way to do this (Boris's
     chunker=TextChunker,
     chunk_size: int = None,
     ontology_file_path: Optional[str] = None,
+    custom_prompt: Optional[str] = None,
 ) -> list[Task]:
     default_tasks = [
         Task(classify_documents),
@@ -308,6 +222,7 @@ async def get_default_tasks(  # TODO: Find out a better way to do this (Boris's
             extract_graph_from_data,
             graph_model=graph_model,
             ontology_adapter=OntologyResolver(ontology_file=ontology_file_path),
+            custom_prompt=custom_prompt,
             task_config={"batch_size": 10},
         ),  # Generate knowledge graphs from the document chunks.
         Task(

{cognee-0.2.3.dev1 → cognee-0.2.4}/cognee/api/v1/cognify/routers/get_cognify_router.py

@@ -37,6 +37,9 @@ class CognifyPayloadDTO(InDTO):
     datasets: Optional[List[str]] = Field(default=None)
     dataset_ids: Optional[List[UUID]] = Field(default=None, examples=[[]])
     run_in_background: Optional[bool] = Field(default=False)
+    custom_prompt: Optional[str] = Field(
+        default=None, description="Custom prompt for entity extraction and graph generation"
+    )
 
 
 def get_cognify_router() -> APIRouter:
@@ -63,6 +66,7 @@ def get_cognify_router() -> APIRouter:
         - **datasets** (Optional[List[str]]): List of dataset names to process. Dataset names are resolved to datasets owned by the authenticated user.
         - **dataset_ids** (Optional[List[UUID]]): List of existing dataset UUIDs to process. UUIDs allow processing of datasets not owned by the user (if permitted).
         - **run_in_background** (Optional[bool]): Whether to execute processing asynchronously. Defaults to False (blocking).
+        - **custom_prompt** (Optional[str]): Custom prompt for entity extraction and graph generation. If provided, this prompt will be used instead of the default prompts for knowledge graph extraction.
 
         ## Response
         - **Blocking execution**: Complete pipeline run information with entity counts, processing duration, and success/failure status
@@ -76,7 +80,8 @@ def get_cognify_router() -> APIRouter:
         ```json
         {
             "datasets": ["research_papers", "documentation"],
-            "run_in_background": false
+            "run_in_background": false,
+            "custom_prompt": "Extract entities focusing on technical concepts and their relationships. Identify key technologies, methodologies, and their interconnections."
         }
         ```
 
@@ -106,7 +111,10 @@ def get_cognify_router() -> APIRouter:
             datasets = payload.dataset_ids if payload.dataset_ids else payload.datasets
 
             cognify_run = await cognee_cognify(
-                datasets, user, run_in_background=payload.run_in_background
+                datasets,
+                user,
+                run_in_background=payload.run_in_background,
+                custom_prompt=payload.custom_prompt,
             )
 
             # If any cognify run errored return JSONResponse with proper error status code
@@ -164,7 +172,7 @@ def get_cognify_router() -> APIRouter:
                     {
                         "pipeline_run_id": str(pipeline_run_info.pipeline_run_id),
                         "status": pipeline_run_info.status,
-                        "payload": await get_formatted_graph_data(pipeline_run.dataset_id, user.id),
+                        "payload": await get_formatted_graph_data(pipeline_run.dataset_id, user),
                     }
                 )
 

{cognee-0.2.3.dev1 → cognee-0.2.4}/cognee/api/v1/datasets/routers/get_datasets_router.py

@@ -284,7 +284,7 @@ def get_datasets_router() -> APIRouter:
         - **500 Internal Server Error**: Error retrieving graph data
         """
 
-        graph_data = await get_formatted_graph_data(dataset_id, user.id)
+        graph_data = await get_formatted_graph_data(dataset_id, user)
 
         return graph_data
 

{cognee-0.2.3.dev1 → cognee-0.2.4}/cognee/api/v1/responses/default_tools.py

@@ -49,6 +49,10 @@ DEFAULT_TOOLS = [
                     "type": "string",
                     "description": "Path to a custom ontology file",
                 },
+                "custom_prompt": {
+                    "type": "string",
+                    "description": "Custom prompt for entity extraction and graph generation. If provided, this prompt will be used instead of the default prompts.",
+                },
             },
             "required": ["text"],
         },

{cognee-0.2.3.dev1 → cognee-0.2.4}/cognee/api/v1/responses/dispatch_function.py

@@ -88,11 +88,16 @@ async def handle_cognify(arguments: Dict[str, Any], user) -> str:
     """Handle cognify function call"""
     text = arguments.get("text")
     ontology_file_path = arguments.get("ontology_file_path")
+    custom_prompt = arguments.get("custom_prompt")
 
     if text:
         await add(data=text, user=user)
 
-    await cognify(user=user, ontology_file_path=ontology_file_path if ontology_file_path else None)
+    await cognify(
+        user=user,
+        ontology_file_path=ontology_file_path if ontology_file_path else None,
+        custom_prompt=custom_prompt,
+    )
 
     return (
         "Text successfully converted into knowledge graph."

{cognee-0.2.3.dev1 → cognee-0.2.4}/cognee/api/v1/responses/models.py

@@ -70,7 +70,7 @@ class ResponseRequest(InDTO):
     tool_choice: Optional[Union[str, Dict[str, Any]]] = "auto"
     user: Optional[str] = None
     temperature: Optional[float] = 1.0
-    max_tokens: Optional[int] = None
+    max_completion_tokens: Optional[int] = None
 
 
 class ToolCallOutput(BaseModel):

{cognee-0.2.3.dev1 → cognee-0.2.4}/cognee/api/v1/search/search.py

@@ -19,6 +19,8 @@ async def search(
     top_k: int = 10,
     node_type: Optional[Type] = None,
     node_name: Optional[List[str]] = None,
+    save_interaction: bool = False,
+    last_k: Optional[int] = None,
 ) -> list:
     """
     Search and query the knowledge graph for insights, information, and connections.
@@ -107,6 +109,8 @@ async def search(
 
         node_name: Filter results to specific named entities (for targeted search).
 
+        save_interaction: Save interaction (query, context, answer connected to triplet endpoints) results into the graph or not
+
     Returns:
         list: Search results in format determined by query_type:
 
@@ -182,6 +186,8 @@ async def search(
         top_k=top_k,
         node_type=node_type,
         node_name=node_name,
+        save_interaction=save_interaction,
+        last_k=last_k,
     )
 
     return filtered_search_results

cognee-0.2.4/cognee/cli/__init__.py

@@ -0,0 +1,10 @@
+from cognee.cli.reference import SupportsCliCommand
+from cognee.cli.exceptions import CliCommandException
+
+DEFAULT_DOCS_URL = "https://docs.cognee.ai"
+
+__all__ = [
+    "SupportsCliCommand",
+    "CliCommandException",
+    "DEFAULT_DOCS_URL",
+]