cognee 0.3.6__py3-none-any.whl → 0.3.7.dev1__py3-none-any.whl

cognee/modules/pipelines/operations/run_tasks_distributed.py

@@ -3,49 +3,97 @@ try:
 except ModuleNotFoundError:
     modal = None
 
+from typing import Any, List, Optional
+from uuid import UUID
+
+from cognee.modules.pipelines.tasks.task import Task
 from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.modules.pipelines.models import (
     PipelineRunStarted,
-    PipelineRunYield,
     PipelineRunCompleted,
+    PipelineRunErrored,
+)
+from cognee.modules.pipelines.operations import (
+    log_pipeline_run_start,
+    log_pipeline_run_complete,
+    log_pipeline_run_error,
 )
-from cognee.modules.pipelines.operations import log_pipeline_run_start, log_pipeline_run_complete
-from cognee.modules.pipelines.utils.generate_pipeline_id import generate_pipeline_id
+from cognee.modules.pipelines.utils import generate_pipeline_id
 from cognee.modules.users.methods import get_default_user
 from cognee.shared.logging_utils import get_logger
-
-from .run_tasks_with_telemetry import run_tasks_with_telemetry
-
+from cognee.modules.users.models import User
+from cognee.modules.pipelines.exceptions import PipelineRunFailedError
+from cognee.tasks.ingestion import resolve_data_directories
+from .run_tasks_data_item import run_tasks_data_item
 
 logger = get_logger("run_tasks_distributed()")
 
-
 if modal:
+    import os
     from distributed.app import app
     from distributed.modal_image import image
 
+    secret_name = os.environ.get("MODAL_SECRET_NAME", "distributed_cognee")
+
     @app.function(
         retries=3,
         image=image,
         timeout=86400,
         max_containers=50,
-        secrets=[modal.Secret.from_name("distributed_cognee")],
+        secrets=[modal.Secret.from_name(secret_name)],
     )
-    async def run_tasks_on_modal(tasks, data_item, user, pipeline_name, context):
-        pipeline_run = run_tasks_with_telemetry(tasks, data_item, user, pipeline_name, context)
-
-        run_info = None
-
-        async for pipeline_run_info in pipeline_run:
-            run_info = pipeline_run_info
+    async def run_tasks_on_modal(
+        data_item,
+        dataset_id: UUID,
+        tasks: List[Task],
+        pipeline_name: str,
+        pipeline_id: str,
+        pipeline_run_id: str,
+        context: Optional[dict],
+        user: User,
+        incremental_loading: bool,
+    ):
+        """
+        Wrapper that runs the run_tasks_data_item function.
+        This is the function/code that runs on modal executor and produces the graph/vector db objects
+        """
+        from cognee.infrastructure.databases.relational import get_relational_engine
+
+        async with get_relational_engine().get_async_session() as session:
+            from cognee.modules.data.models import Dataset
+
+            dataset = await session.get(Dataset, dataset_id)
+
+        result = await run_tasks_data_item(
+            data_item=data_item,
+            dataset=dataset,
+            tasks=tasks,
+            pipeline_name=pipeline_name,
+            pipeline_id=pipeline_id,
+            pipeline_run_id=pipeline_run_id,
+            context=context,
+            user=user,
+            incremental_loading=incremental_loading,
+        )
 
-        return run_info
+        return result
 
 
-async def run_tasks_distributed(tasks, dataset_id, data, user, pipeline_name, context):
+async def run_tasks_distributed(
+    tasks: List[Task],
+    dataset_id: UUID,
+    data: List[Any] = None,
+    user: User = None,
+    pipeline_name: str = "unknown_pipeline",
+    context: dict = None,
+    incremental_loading: bool = False,
+    data_per_batch: int = 20,
+):
     if not user:
         user = await get_default_user()
 
+    # Get dataset object
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         from cognee.modules.data.models import Dataset
@@ -53,9 +101,7 @@ async def run_tasks_distributed(tasks, dataset_id, data, user, pipeline_name, co
         dataset = await session.get(Dataset, dataset_id)
 
     pipeline_id = generate_pipeline_id(user.id, dataset.id, pipeline_name)
-
     pipeline_run = await log_pipeline_run_start(pipeline_id, pipeline_name, dataset_id, data)
-
     pipeline_run_id = pipeline_run.pipeline_run_id
 
     yield PipelineRunStarted(
@@ -65,30 +111,67 @@ async def run_tasks_distributed(tasks, dataset_id, data, user, pipeline_name, co
         payload=data,
     )
 
-    data_count = len(data) if isinstance(data, list) else 1
+    try:
+        if not isinstance(data, list):
+            data = [data]
+
+        data = await resolve_data_directories(data)
+
+        number_of_data_items = len(data) if isinstance(data, list) else 1
+
+        data_item_tasks = [
+            data,
+            [dataset.id] * number_of_data_items,
+            [tasks] * number_of_data_items,
+            [pipeline_name] * number_of_data_items,
+            [pipeline_id] * number_of_data_items,
+            [pipeline_run_id] * number_of_data_items,
+            [context] * number_of_data_items,
+            [user] * number_of_data_items,
+            [incremental_loading] * number_of_data_items,
+        ]
+
+        results = []
+        async for result in run_tasks_on_modal.map.aio(*data_item_tasks):
+            if not result:
+                continue
+            results.append(result)
+
+        # Remove skipped results
+        results = [r for r in results if r]
+
+        # If any data item failed, raise PipelineRunFailedError
+        errored = [
+            r
+            for r in results
+            if r and r.get("run_info") and isinstance(r["run_info"], PipelineRunErrored)
+        ]
+        if errored:
+            raise PipelineRunFailedError("Pipeline run failed. Data item could not be processed.")
+
+        await log_pipeline_run_complete(
+            pipeline_run_id, pipeline_id, pipeline_name, dataset_id, data
+        )
 
-    arguments = [
-        [tasks] * data_count,
-        [[data_item] for data_item in data[:data_count]] if data_count > 1 else [data],
-        [user] * data_count,
-        [pipeline_name] * data_count,
-        [context] * data_count,
-    ]
+        yield PipelineRunCompleted(
+            pipeline_run_id=pipeline_run_id,
+            dataset_id=dataset.id,
+            dataset_name=dataset.name,
+            data_ingestion_info=results,
+        )
 
-    async for result in run_tasks_on_modal.map.aio(*arguments):
-        logger.info(f"Received result: {result}")
+    except Exception as error:
+        await log_pipeline_run_error(
+            pipeline_run_id, pipeline_id, pipeline_name, dataset_id, data, error
+        )
 
-        yield PipelineRunYield(
+        yield PipelineRunErrored(
             pipeline_run_id=pipeline_run_id,
+            payload=repr(error),
             dataset_id=dataset.id,
             dataset_name=dataset.name,
-            payload=result,
+            data_ingestion_info=locals().get("results"),
         )
 
-    await log_pipeline_run_complete(pipeline_run_id, pipeline_id, pipeline_name, dataset_id, data)
-
-    yield PipelineRunCompleted(
-        pipeline_run_id=pipeline_run_id,
-        dataset_id=dataset.id,
-        dataset_name=dataset.name,
-    )
+        if not isinstance(error, PipelineRunFailedError):
+            raise

cognee/modules/pipelines/operations/run_tasks_with_telemetry.py

@@ -4,6 +4,7 @@ from cognee.modules.settings import get_current_settings
 from cognee.modules.users.models import User
 from cognee.shared.logging_utils import get_logger
 from cognee.shared.utils import send_telemetry
+from cognee import __version__ as cognee_version
 
 from .run_tasks_base import run_tasks_base
 from ..tasks.task import Task
@@ -26,6 +27,8 @@ async def run_tasks_with_telemetry(
             user.id,
             additional_properties={
                 "pipeline_name": str(pipeline_name),
+                "cognee_version": cognee_version,
+                "tenant_id": str(user.tenant_id) if user.tenant_id else "Single User Tenant",
             }
             | config,
         )
@@ -39,7 +42,10 @@ async def run_tasks_with_telemetry(
             user.id,
             additional_properties={
                 "pipeline_name": str(pipeline_name),
-            },
+                "cognee_version": cognee_version,
+                "tenant_id": str(user.tenant_id) if user.tenant_id else "Single User Tenant",
+            }
+            | config,
         )
     except Exception as error:
         logger.error(
@@ -53,6 +59,8 @@ async def run_tasks_with_telemetry(
             user.id,
             additional_properties={
                 "pipeline_name": str(pipeline_name),
+                "cognee_version": cognee_version,
+                "tenant_id": str(user.tenant_id) if user.tenant_id else "Single User Tenant",
             }
             | config,
         )

cognee/modules/retrieval/EntityCompletionRetriever.py

@@ -1,10 +1,17 @@
+import asyncio
 from typing import Any, Optional, List
 from cognee.shared.logging_utils import get_logger
 
 from cognee.infrastructure.entities.BaseEntityExtractor import BaseEntityExtractor
 from cognee.infrastructure.context.BaseContextProvider import BaseContextProvider
 from cognee.modules.retrieval.base_retriever import BaseRetriever
-from cognee.modules.retrieval.utils.completion import generate_completion
+from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
+from cognee.modules.retrieval.utils.session_cache import (
+    save_conversation_history,
+    get_conversation_history,
+)
+from cognee.context_global_variables import session_user
+from cognee.infrastructure.databases.cache.config import CacheConfig
 
 
 logger = get_logger("entity_completion_retriever")
@@ -77,7 +84,9 @@ class EntityCompletionRetriever(BaseRetriever):
             logger.error(f"Context retrieval failed: {str(e)}")
             return None
 
-    async def get_completion(self, query: str, context: Optional[Any] = None) -> List[str]:
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None
+    ) -> List[str]:
         """
         Generate completion using provided context or fetch new context.
 
@@ -91,6 +100,8 @@ class EntityCompletionRetriever(BaseRetriever):
             - query (str): The query string for which completion is being generated.
             - context (Optional[Any]): Optional context to be used for generating completion;
               fetched if not provided. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
 
         Returns:
         --------
@@ -105,12 +116,41 @@ class EntityCompletionRetriever(BaseRetriever):
             if context is None:
                 return ["No relevant entities found for the query."]
 
-            completion = await generate_completion(
-                query=query,
-                context=context,
-                user_prompt_path=self.user_prompt_path,
-                system_prompt_path=self.system_prompt_path,
-            )
+            # Check if we need to generate context summary for caching
+            cache_config = CacheConfig()
+            user = session_user.get()
+            user_id = getattr(user, "id", None)
+            session_save = user_id and cache_config.caching
+
+            if session_save:
+                conversation_history = await get_conversation_history(session_id=session_id)
+
+                context_summary, completion = await asyncio.gather(
+                    summarize_text(str(context)),
+                    generate_completion(
+                        query=query,
+                        context=context,
+                        user_prompt_path=self.user_prompt_path,
+                        system_prompt_path=self.system_prompt_path,
+                        conversation_history=conversation_history,
+                    ),
+                )
+            else:
+                completion = await generate_completion(
+                    query=query,
+                    context=context,
+                    user_prompt_path=self.user_prompt_path,
+                    system_prompt_path=self.system_prompt_path,
+                )
+
+            if session_save:
+                await save_conversation_history(
+                    query=query,
+                    context_summary=context_summary,
+                    answer=completion,
+                    session_id=session_id,
+                )
+
             return [completion]
 
         except Exception as e:

cognee/modules/retrieval/base_graph_retriever.py

@@ -13,6 +13,8 @@ class BaseGraphRetriever(ABC):
         pass
 
     @abstractmethod
-    async def get_completion(self, query: str, context: Optional[List[Edge]] = None) -> str:
+    async def get_completion(
+        self, query: str, context: Optional[List[Edge]] = None, session_id: Optional[str] = None
+    ) -> str:
         """Generates a response using the query and optional context (triplets)."""
         pass

cognee/modules/retrieval/base_retriever.py

@@ -11,6 +11,8 @@ class BaseRetriever(ABC):
         pass
 
     @abstractmethod
-    async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None
+    ) -> Any:
         """Generates a response using the query and optional context."""
         pass

cognee/modules/retrieval/chunks_retriever.py

@@ -61,7 +61,9 @@ class ChunksRetriever(BaseRetriever):
         logger.info(f"Returning {len(chunk_payloads)} chunk payloads")
         return chunk_payloads
 
-    async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None
+    ) -> Any:
         """
         Generates a completion using document chunks context.
 
@@ -74,6 +76,8 @@ class ChunksRetriever(BaseRetriever):
             - query (str): The query string to be used for generating a completion.
             - context (Optional[Any]): Optional pre-fetched context to use for generating the
               completion; if None, it retrieves the context for the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
 
         Returns:
         --------

cognee/modules/retrieval/code_retriever.py

@@ -207,8 +207,26 @@ class CodeRetriever(BaseRetriever):
         logger.info(f"Returning {len(result)} code file contexts")
         return result
 
-    async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
-        """Returns the code files context."""
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None
+    ) -> Any:
+        """
+        Returns the code files context.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string to retrieve code context for.
+            - context (Optional[Any]): Optional pre-fetched context; if None, it retrieves
+              the context for the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
+
+        Returns:
+        --------
+
+            - Any: The code files context, either provided or retrieved.
+        """
         if context is None:
             context = await self.get_context(query)
         return context

cognee/modules/retrieval/completion_retriever.py

@@ -1,11 +1,18 @@
+import asyncio
 from typing import Any, Optional
 
 from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.databases.vector import get_vector_engine
-from cognee.modules.retrieval.utils.completion import generate_completion
+from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
+from cognee.modules.retrieval.utils.session_cache import (
+    save_conversation_history,
+    get_conversation_history,
+)
 from cognee.modules.retrieval.base_retriever import BaseRetriever
 from cognee.modules.retrieval.exceptions.exceptions import NoDataError
 from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
+from cognee.context_global_variables import session_user
+from cognee.infrastructure.databases.cache.config import CacheConfig
 
 logger = get_logger("CompletionRetriever")
 
@@ -67,7 +74,9 @@ class CompletionRetriever(BaseRetriever):
             logger.error("DocumentChunk_text collection not found")
             raise NoDataError("No data found in the system, please add data first.") from error
 
-    async def get_completion(self, query: str, context: Optional[Any] = None) -> str:
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None
+    ) -> str:
         """
         Generates an LLM completion using the context.
 
@@ -80,6 +89,8 @@ class CompletionRetriever(BaseRetriever):
             - query (str): The query string to be used for generating a completion.
             - context (Optional[Any]): Optional pre-fetched context to use for generating the
               completion; if None, it retrieves the context for the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
 
         Returns:
         --------
@@ -89,11 +100,41 @@ class CompletionRetriever(BaseRetriever):
         if context is None:
             context = await self.get_context(query)
 
-        completion = await generate_completion(
-            query=query,
-            context=context,
-            user_prompt_path=self.user_prompt_path,
-            system_prompt_path=self.system_prompt_path,
-            system_prompt=self.system_prompt,
-        )
+        # Check if we need to generate context summary for caching
+        cache_config = CacheConfig()
+        user = session_user.get()
+        user_id = getattr(user, "id", None)
+        session_save = user_id and cache_config.caching
+
+        if session_save:
+            conversation_history = await get_conversation_history(session_id=session_id)
+
+            context_summary, completion = await asyncio.gather(
+                summarize_text(context),
+                generate_completion(
+                    query=query,
+                    context=context,
+                    user_prompt_path=self.user_prompt_path,
+                    system_prompt_path=self.system_prompt_path,
+                    system_prompt=self.system_prompt,
+                    conversation_history=conversation_history,
+                ),
+            )
+        else:
+            completion = await generate_completion(
+                query=query,
+                context=context,
+                user_prompt_path=self.user_prompt_path,
+                system_prompt_path=self.system_prompt_path,
+                system_prompt=self.system_prompt,
+            )
+
+        if session_save:
+            await save_conversation_history(
+                query=query,
+                context_summary=context_summary,
+                answer=completion,
+                session_id=session_id,
+            )
+
         return completion

cognee/modules/retrieval/cypher_search_retriever.py

@@ -44,13 +44,21 @@ class CypherSearchRetriever(BaseRetriever):
         """
         try:
             graph_engine = await get_graph_engine()
+            is_empty = await graph_engine.is_empty()
+
+            if is_empty:
+                logger.warning("Search attempt on an empty knowledge graph")
+                return []
+
             result = await graph_engine.query(query)
         except Exception as e:
             logger.error("Failed to execture cypher search retrieval: %s", str(e))
             raise CypherSearchError() from e
         return result
 
-    async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None
+    ) -> Any:
         """
         Returns the graph connections context.
 
@@ -62,6 +70,8 @@ class CypherSearchRetriever(BaseRetriever):
             - query (str): The query to retrieve context.
             - context (Optional[Any]): Optional context to use, otherwise fetched using the
               query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
 
         Returns:
         --------

cognee/modules/retrieval/graph_completion_context_extension_retriever.py

@@ -1,8 +1,15 @@
+import asyncio
 from typing import Optional, List, Type
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
-from cognee.modules.retrieval.utils.completion import generate_completion
+from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
+from cognee.modules.retrieval.utils.session_cache import (
+    save_conversation_history,
+    get_conversation_history,
+)
+from cognee.context_global_variables import session_user
+from cognee.infrastructure.databases.cache.config import CacheConfig
 
 logger = get_logger()
 
@@ -47,6 +54,7 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
         self,
         query: str,
         context: Optional[List[Edge]] = None,
+        session_id: Optional[str] = None,
         context_extension_rounds=4,
     ) -> List[str]:
         """
@@ -64,6 +72,8 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
             - query (str): The input query for which the completion is generated.
             - context (Optional[Any]): The existing context to use for enhancing the query; if
               None, it will be initialized from triplets generated for the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
             - context_extension_rounds: The maximum number of rounds to extend the context with
               new triplets before halting. (default 4)
 
@@ -115,17 +125,46 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
 
             round_idx += 1
 
-        completion = await generate_completion(
-            query=query,
-            context=context_text,
-            user_prompt_path=self.user_prompt_path,
-            system_prompt_path=self.system_prompt_path,
-            system_prompt=self.system_prompt,
-        )
+        # Check if we need to generate context summary for caching
+        cache_config = CacheConfig()
+        user = session_user.get()
+        user_id = getattr(user, "id", None)
+        session_save = user_id and cache_config.caching
+
+        if session_save:
+            conversation_history = await get_conversation_history(session_id=session_id)
+
+            context_summary, completion = await asyncio.gather(
+                summarize_text(context_text),
+                generate_completion(
+                    query=query,
+                    context=context_text,
+                    user_prompt_path=self.user_prompt_path,
+                    system_prompt_path=self.system_prompt_path,
+                    system_prompt=self.system_prompt,
+                    conversation_history=conversation_history,
+                ),
+            )
+        else:
+            completion = await generate_completion(
+                query=query,
+                context=context_text,
+                user_prompt_path=self.user_prompt_path,
+                system_prompt_path=self.system_prompt_path,
+                system_prompt=self.system_prompt,
+            )
 
         if self.save_interaction and context_text and triplets and completion:
             await self.save_qa(
                 question=query, answer=completion, context=context_text, triplets=triplets
             )
 
+        if session_save:
+            await save_conversation_history(
+                query=query,
+                context_summary=context_summary,
+                answer=completion,
+                session_id=session_id,
+            )
+
         return [completion]