cognee 0.2.3.dev1__py3-none-any.whl → 0.3.0__py3-none-any.whl

cognee/modules/pipelines/operations/pipeline.py

@@ -2,41 +2,23 @@ import asyncio
 from uuid import UUID
 from typing import Union
 
+from cognee.modules.pipelines.layers.setup_and_check_environment import (
+    setup_and_check_environment,
+)
+
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.data.methods.get_dataset_data import get_dataset_data
 from cognee.modules.data.models import Data, Dataset
 from cognee.modules.pipelines.operations.run_tasks import run_tasks
-from cognee.modules.pipelines.models import PipelineRunStatus
-from cognee.modules.pipelines.utils import generate_pipeline_id
-from cognee.modules.pipelines.operations.get_pipeline_status import get_pipeline_status
-from cognee.modules.pipelines.methods import get_pipeline_run_by_dataset
-
+from cognee.modules.pipelines.layers import validate_pipeline_tasks
 from cognee.modules.pipelines.tasks.task import Task
-from cognee.modules.users.methods import get_default_user
 from cognee.modules.users.models import User
-from cognee.modules.pipelines.operations import log_pipeline_run_initiated
 from cognee.context_global_variables import set_database_global_context_variables
-from cognee.modules.data.exceptions import DatasetNotFoundError
-from cognee.modules.data.methods import (
-    get_authorized_existing_datasets,
-    load_or_create_datasets,
-    check_dataset_name,
+from cognee.modules.pipelines.layers.resolve_authorized_user_datasets import (
+    resolve_authorized_user_datasets,
 )
-
-from cognee.modules.pipelines.models.PipelineRunInfo import (
-    PipelineRunCompleted,
-    PipelineRunStarted,
-)
-
-from cognee.infrastructure.databases.relational import (
-    create_db_and_tables as create_relational_db_and_tables,
-)
-from cognee.infrastructure.databases.vector.pgvector import (
-    create_db_and_tables as create_pgvector_db_and_tables,
-)
-from cognee.context_global_variables import (
-    graph_db_config as context_graph_db_config,
-    vector_db_config as context_vector_db_config,
+from cognee.modules.pipelines.layers.check_pipeline_run_qualification import (
+    check_pipeline_run_qualification,
 )
 
 logger = get_logger("cognee.pipeline")
@@ -44,7 +26,7 @@ logger = get_logger("cognee.pipeline")
 update_status_lock = asyncio.Lock()
 
 
-async def cognee_pipeline(
+async def run_pipeline(
     tasks: list[Task],
     data=None,
     datasets: Union[str, list[str], list[UUID]] = None,
@@ -54,56 +36,13 @@ async def cognee_pipeline(
     graph_db_config: dict = None,
     incremental_loading: bool = False,
 ):
-    # Note: These context variables allow different value assignment for databases in Cognee
-    #       per async task, thread, process and etc.
-    if vector_db_config:
-        context_vector_db_config.set(vector_db_config)
-    if graph_db_config:
-        context_graph_db_config.set(graph_db_config)
-
-    # Create tables for databases
-    await create_relational_db_and_tables()
-    await create_pgvector_db_and_tables()
-
-    # Initialize first_run attribute if it doesn't exist
-    if not hasattr(cognee_pipeline, "first_run"):
-        cognee_pipeline.first_run = True
-
-    if cognee_pipeline.first_run:
-        from cognee.infrastructure.llm.utils import (
-            test_llm_connection,
-            test_embedding_connection,
-        )
-
-        # Test LLM and Embedding configuration once before running Cognee
-        await test_llm_connection()
-        await test_embedding_connection()
-        cognee_pipeline.first_run = False  # Update flag after first run
+    validate_pipeline_tasks(tasks)
+    await setup_and_check_environment(vector_db_config, graph_db_config)
 
-    # If no user is provided use default user
-    if user is None:
-        user = await get_default_user()
+    user, authorized_datasets = await resolve_authorized_user_datasets(datasets, user)
 
-    # Convert datasets to list
-    if isinstance(datasets, str) or isinstance(datasets, UUID):
-        datasets = [datasets]
-
-    # Get datasets user wants write permissions for (verify user has permissions if datasets are provided as well)
-    # NOTE: If a user wants to write to a dataset he does not own it must be provided through UUID
-    existing_datasets = await get_authorized_existing_datasets(datasets, "write", user)
-
-    if not datasets:
-        # Get datasets from database if none sent.
-        datasets = existing_datasets
-    else:
-        # If dataset matches an existing Dataset (by name or id), reuse it. Otherwise, create a new Dataset.
-        datasets = await load_or_create_datasets(datasets, existing_datasets, user)
-
-    if not datasets:
-        raise DatasetNotFoundError("There are no datasets to work with.")
-
-    for dataset in datasets:
-        async for run_info in run_pipeline(
+    for dataset in authorized_datasets:
+        async for run_info in run_pipeline_per_dataset(
             dataset=dataset,
             user=user,
             tasks=tasks,
@@ -115,7 +54,7 @@ async def cognee_pipeline(
             yield run_info
 
 
-async def run_pipeline(
+async def run_pipeline_per_dataset(
     dataset: Dataset,
     user: User,
     tasks: list[Task],
@@ -124,74 +63,21 @@ async def run_pipeline(
     context: dict = None,
     incremental_loading=False,
 ):
-    check_dataset_name(dataset.name)
-
     # Will only be used if ENABLE_BACKEND_ACCESS_CONTROL is set to True
     await set_database_global_context_variables(dataset.id, dataset.owner_id)
 
-    # Ugly hack, but no easier way to do this.
-    if pipeline_name == "add_pipeline":
-        pipeline_id = generate_pipeline_id(user.id, dataset.id, pipeline_name)
-        # Refresh the add pipeline status so data is added to a dataset.
-        # Without this the app_pipeline status will be DATASET_PROCESSING_COMPLETED and will skip the execution.
-
-        await log_pipeline_run_initiated(
-            pipeline_id=pipeline_id,
-            pipeline_name="add_pipeline",
-            dataset_id=dataset.id,
-        )
-
-        # Refresh the cognify pipeline status after we add new files.
-        # Without this the cognify_pipeline status will be DATASET_PROCESSING_COMPLETED and will skip the execution.
-        await log_pipeline_run_initiated(
-            pipeline_id=pipeline_id,
-            pipeline_name="cognify_pipeline",
-            dataset_id=dataset.id,
-        )
-
-    dataset_id = dataset.id
-
     if not data:
-        data: list[Data] = await get_dataset_data(dataset_id=dataset_id)
-
-    # async with update_status_lock: TODO: Add UI lock to prevent multiple backend requests
-    if isinstance(dataset, Dataset):
-        task_status = await get_pipeline_status([dataset_id], pipeline_name)
-    else:
-        task_status = [
-            PipelineRunStatus.DATASET_PROCESSING_COMPLETED
-        ]  # TODO: this is a random assignment, find permanent solution
-
-    if str(dataset_id) in task_status:
-        if task_status[str(dataset_id)] == PipelineRunStatus.DATASET_PROCESSING_STARTED:
-            logger.info("Dataset %s is already being processed.", dataset_id)
-            pipeline_run = await get_pipeline_run_by_dataset(dataset_id, pipeline_name)
-            yield PipelineRunStarted(
-                pipeline_run_id=pipeline_run.pipeline_run_id,
-                dataset_id=dataset.id,
-                dataset_name=dataset.name,
-                payload=data,
-            )
-            return
-        elif task_status[str(dataset_id)] == PipelineRunStatus.DATASET_PROCESSING_COMPLETED:
-            logger.info("Dataset %s is already processed.", dataset_id)
-            pipeline_run = await get_pipeline_run_by_dataset(dataset_id, pipeline_name)
-            yield PipelineRunCompleted(
-                pipeline_run_id=pipeline_run.pipeline_run_id,
-                dataset_id=dataset.id,
-                dataset_name=dataset.name,
-            )
-            return
-
-    if not isinstance(tasks, list):
-        raise ValueError("Tasks must be a list")
+        data: list[Data] = await get_dataset_data(dataset_id=dataset.id)
 
-    for task in tasks:
-        if not isinstance(task, Task):
-            raise ValueError(f"Task {task} is not an instance of Task")
+    process_pipeline_status = await check_pipeline_run_qualification(dataset, data, pipeline_name)
+    if process_pipeline_status:
+        # If pipeline was already processed or is currently being processed
+        # return status information to async generator and finish execution
+        yield process_pipeline_status
+        return
 
     pipeline_run = run_tasks(
-        tasks, dataset_id, data, user, pipeline_name, context, incremental_loading
+        tasks, dataset.id, data, user, pipeline_name, context, incremental_loading
     )
 
     async for pipeline_run_info in pipeline_run:

cognee/modules/pipelines/operations/run_tasks.py

@@ -266,48 +266,24 @@ async def run_tasks(
         if incremental_loading:
             data = await resolve_data_directories(data)
 
-        # TODO: Return to using async.gather for data items after Cognee release
-        # # Create async tasks per data item that will run the pipeline for the data item
-        # data_item_tasks = [
-        #     asyncio.create_task(
-        #         _run_tasks_data_item(
-        #             data_item,
-        #             dataset,
-        #             tasks,
-        #             pipeline_name,
-        #             pipeline_id,
-        #             pipeline_run_id,
-        #             context,
-        #             user,
-        #             incremental_loading,
-        #         )
-        #     )
-        #     for data_item in data
-        # ]
-        # results = await asyncio.gather(*data_item_tasks)
-        # # Remove skipped data items from results
-        # results = [result for result in results if result]
-
-        ### TEMP sync data item handling
-        results = []
-        # Run the pipeline for each data_item sequentially, one after the other
-        for data_item in data:
-            result = await _run_tasks_data_item(
-                data_item,
-                dataset,
-                tasks,
-                pipeline_name,
-                pipeline_id,
-                pipeline_run_id,
-                context,
-                user,
-                incremental_loading,
+        # Create async tasks per data item that will run the pipeline for the data item
+        data_item_tasks = [
+            asyncio.create_task(
+                _run_tasks_data_item(
+                    data_item,
+                    dataset,
+                    tasks,
+                    pipeline_name,
+                    pipeline_id,
+                    pipeline_run_id,
+                    context,
+                    user,
+                    incremental_loading,
+                )
             )
-
-            # Skip items that returned a false-y value
-            if result:
-                results.append(result)
-        ### END
+            for data_item in data
+        ]
+        results = await asyncio.gather(*data_item_tasks)
 
         # Remove skipped data items from results
         results = [result for result in results if result]

cognee/modules/retrieval/base_feedback.py

@@ -0,0 +1,11 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class BaseFeedback(ABC):
+    """Base class for all user feedback operations."""
+
+    @abstractmethod
+    async def add_feedback(self, feedback_text: str) -> Any:
+        """Add user feedback to the system."""
+        pass

cognee/modules/retrieval/base_graph_retriever.py

@@ -0,0 +1,18 @@
+from typing import List, Optional
+from abc import ABC, abstractmethod
+
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+
+
+class BaseGraphRetriever(ABC):
+    """Base class for all graph based retrievers."""
+
+    @abstractmethod
+    async def get_context(self, query: str) -> List[Edge]:
+        """Retrieves triplets based on the query."""
+        pass
+
+    @abstractmethod
+    async def get_completion(self, query: str, context: Optional[List[Edge]] = None) -> str:
+        """Generates a response using the query and optional context (triplets)."""
+        pass

cognee/modules/retrieval/base_retriever.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Optional, Callable
+from typing import Any, Optional
 
 
 class BaseRetriever(ABC):

cognee/modules/retrieval/code_retriever.py

@@ -94,7 +94,15 @@ class CodeRetriever(BaseRetriever):
                         {"id": res.id, "score": res.score, "payload": res.payload}
                     )
 
+            existing_collection = []
             for collection in self.classes_and_functions_collections:
+                if await vector_engine.has_collection(collection):
+                    existing_collection.append(collection)
+
+            if not existing_collection:
+                raise RuntimeError("No collection found for code retriever")
+
+            for collection in existing_collection:
                 logger.debug(f"Searching {collection} collection with general query")
                 search_results_code = await vector_engine.search(
                     collection, query, limit=self.top_k

cognee/modules/retrieval/coding_rules_retriever.py

@@ -0,0 +1,31 @@
+import asyncio
+from functools import reduce
+from typing import List, Optional
+from cognee.shared.logging_utils import get_logger
+from cognee.tasks.codingagents.coding_rule_associations import get_existing_rules
+
+logger = get_logger("CodingRulesRetriever")
+
+
+class CodingRulesRetriever:
+    """Retriever for handling codeing rule based searches."""
+
+    def __init__(self, rules_nodeset_name: Optional[List[str]] = None):
+        if isinstance(rules_nodeset_name, list):
+            if not rules_nodeset_name:
+                # If there is no provided nodeset set to coding_agent_rules
+                rules_nodeset_name = ["coding_agent_rules"]
+
+        self.rules_nodeset_name = rules_nodeset_name
+        """Initialize retriever with search parameters."""
+
+    async def get_existing_rules(self, query_text):
+        if self.rules_nodeset_name:
+            rules_list = await asyncio.gather(
+                *[
+                    get_existing_rules(rules_nodeset_name=nodeset)
+                    for nodeset in self.rules_nodeset_name
+                ]
+            )
+
+            return reduce(lambda x, y: x + y, rules_list, [])

cognee/modules/retrieval/completion_retriever.py

@@ -23,12 +23,14 @@ class CompletionRetriever(BaseRetriever):
         self,
         user_prompt_path: str = "context_for_question.txt",
         system_prompt_path: str = "answer_simple_question.txt",
+        system_prompt: Optional[str] = None,
         top_k: Optional[int] = 1,
     ):
         """Initialize retriever with optional custom prompt paths."""
         self.user_prompt_path = user_prompt_path
         self.system_prompt_path = system_prompt_path
         self.top_k = top_k if top_k is not None else 1
+        self.system_prompt = system_prompt
 
     async def get_context(self, query: str) -> str:
         """
@@ -65,7 +67,7 @@ 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) -> Any:
+    async def get_completion(self, query: str, context: Optional[Any] = None) -> str:
         """
         Generates an LLM completion using the context.
 
@@ -88,6 +90,10 @@ class CompletionRetriever(BaseRetriever):
             context = await self.get_context(query)
 
         completion = await generate_completion(
-            query, context, self.user_prompt_path, self.system_prompt_path
+            query=query,
+            context=context,
+            user_prompt_path=self.user_prompt_path,
+            system_prompt_path=self.system_prompt_path,
+            system_prompt=self.system_prompt,
         )
-        return [completion]
+        return completion

cognee/modules/retrieval/context_providers/TripletSearchContextProvider.py

@@ -49,6 +49,7 @@ class TripletSearchContextProvider(BaseContextProvider):
         tasks = [
             brute_force_triplet_search(
                 query=f"{entity_text} {query}",
+                user=user,
                 top_k=self.top_k,
                 collections=self.collections,
                 properties_to_project=self.properties_to_project,

cognee/modules/retrieval/cypher_search_retriever.py

@@ -1,6 +1,5 @@
 from typing import Any, Optional
 from cognee.infrastructure.databases.graph import get_graph_engine
-from cognee.infrastructure.databases.graph.networkx.adapter import NetworkXAdapter
 from cognee.modules.retrieval.base_retriever import BaseRetriever
 from cognee.modules.retrieval.utils.completion import generate_completion
 from cognee.modules.retrieval.exceptions import SearchTypeNotSupported, CypherSearchError
@@ -31,8 +30,7 @@ class CypherSearchRetriever(BaseRetriever):
         """
         Retrieves relevant context using a cypher query.
 
-        If the graph engine is an instance of NetworkXAdapter, raises SearchTypeNotSupported. If
-        any error occurs during execution, logs the error and raises CypherSearchError.
+        If any error occurs during execution, logs the error and raises CypherSearchError.
 
         Parameters:
         -----------
@@ -46,12 +44,6 @@ class CypherSearchRetriever(BaseRetriever):
         """
         try:
             graph_engine = await get_graph_engine()
-
-            if isinstance(graph_engine, NetworkXAdapter):
-                raise SearchTypeNotSupported(
-                    "CYPHER search type not supported for NetworkXAdapter."
-                )
-
             result = await graph_engine.query(query)
         except Exception as e:
             logger.error("Failed to execture cypher search retrieval: %s", str(e))

cognee/modules/retrieval/graph_completion_context_extension_retriever.py

@@ -1,4 +1,5 @@
-from typing import Any, Optional, List, Type
+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
@@ -26,9 +27,11 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
         self,
         user_prompt_path: str = "graph_context_for_question.txt",
         system_prompt_path: str = "answer_simple_question.txt",
+        system_prompt: Optional[str] = None,
         top_k: Optional[int] = 5,
         node_type: Optional[Type] = None,
         node_name: Optional[List[str]] = None,
+        save_interaction: bool = False,
     ):
         super().__init__(
             user_prompt_path=user_prompt_path,
@@ -36,11 +39,16 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
             top_k=top_k,
             node_type=node_type,
             node_name=node_name,
+            save_interaction=save_interaction,
+            system_prompt=system_prompt,
         )
 
     async def get_completion(
-        self, query: str, context: Optional[Any] = None, context_extension_rounds=4
-    ) -> List[str]:
+        self,
+        query: str,
+        context: Optional[List[Edge]] = None,
+        context_extension_rounds=4,
+    ) -> str:
         """
         Extends the context for a given query by retrieving related triplets and generating new
         completions based on them.
@@ -65,11 +73,12 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
             - List[str]: A list containing the generated answer based on the query and the
               extended context.
         """
-        triplets = []
+        triplets = context
+
+        if triplets is None:
+            triplets = await self.get_context(query)
 
-        if context is None:
-            triplets += await self.get_triplets(query)
-            context = await self.resolve_edges_to_text(triplets)
+        context_text = await self.resolve_edges_to_text(triplets)
 
         round_idx = 1
 
@@ -81,14 +90,15 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
             )
             completion = await generate_completion(
                 query=query,
-                context=context,
+                context=context_text,
                 user_prompt_path=self.user_prompt_path,
                 system_prompt_path=self.system_prompt_path,
+                system_prompt=self.system_prompt,
             )
 
-            triplets += await self.get_triplets(completion)
+            triplets += await self.get_context(completion)
             triplets = list(set(triplets))
-            context = await self.resolve_edges_to_text(triplets)
+            context_text = await self.resolve_edges_to_text(triplets)
 
             num_triplets = len(triplets)
 
@@ -105,11 +115,17 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
 
             round_idx += 1
 
-        answer = await generate_completion(
+        completion = await generate_completion(
             query=query,
-            context=context,
+            context=context_text,
             user_prompt_path=self.user_prompt_path,
             system_prompt_path=self.system_prompt_path,
+            system_prompt=self.system_prompt,
         )
 
-        return [answer]
+        if self.save_interaction and context_text and triplets and completion:
+            await self.save_qa(
+                question=query, answer=completion, context=context_text, triplets=triplets
+            )
+
+        return completion

cognee/modules/retrieval/graph_completion_cot_retriever.py

@@ -1,4 +1,5 @@
-from typing import Any, Optional, List, Type
+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
@@ -32,16 +33,20 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
         validation_system_prompt_path: str = "cot_validation_system_prompt.txt",
         followup_system_prompt_path: str = "cot_followup_system_prompt.txt",
         followup_user_prompt_path: str = "cot_followup_user_prompt.txt",
+        system_prompt: Optional[str] = None,
         top_k: Optional[int] = 5,
         node_type: Optional[Type] = None,
         node_name: Optional[List[str]] = None,
+        save_interaction: bool = False,
     ):
         super().__init__(
             user_prompt_path=user_prompt_path,
             system_prompt_path=system_prompt_path,
+            system_prompt=system_prompt,
             top_k=top_k,
             node_type=node_type,
             node_name=node_name,
+            save_interaction=save_interaction,
         )
         self.validation_system_prompt_path = validation_system_prompt_path
         self.validation_user_prompt_path = validation_user_prompt_path
@@ -49,8 +54,11 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
         self.followup_user_prompt_path = followup_user_prompt_path
 
     async def get_completion(
-        self, query: str, context: Optional[Any] = None, max_iter=4
-    ) -> List[str]:
+        self,
+        query: str,
+        context: Optional[List[Edge]] = None,
+        max_iter=4,
+    ) -> str:
         """
         Generate completion responses based on a user query and contextual information.
 
@@ -75,25 +83,29 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
         """
         followup_question = ""
         triplets = []
-        answer = [""]
+        completion = ""
 
         for round_idx in range(max_iter + 1):
             if round_idx == 0:
                 if context is None:
-                    context = await self.get_context(query)
+                    triplets = await self.get_context(query)
+                    context_text = await self.resolve_edges_to_text(triplets)
+                else:
+                    context_text = await self.resolve_edges_to_text(context)
             else:
-                triplets += await self.get_triplets(followup_question)
-                context = await self.resolve_edges_to_text(list(set(triplets)))
+                triplets += await self.get_context(followup_question)
+                context_text = await self.resolve_edges_to_text(list(set(triplets)))
 
-            answer = await generate_completion(
+            completion = await generate_completion(
                 query=query,
-                context=context,
+                context=context_text,
                 user_prompt_path=self.user_prompt_path,
                 system_prompt_path=self.system_prompt_path,
+                system_prompt=self.system_prompt,
             )
-            logger.info(f"Chain-of-thought: round {round_idx} - answer: {answer}")
+            logger.info(f"Chain-of-thought: round {round_idx} - answer: {completion}")
             if round_idx < max_iter:
-                valid_args = {"query": query, "answer": answer, "context": context}
+                valid_args = {"query": query, "answer": completion, "context": context_text}
                 valid_user_prompt = LLMGateway.render_prompt(
                     filename=self.validation_user_prompt_path, context=valid_args
                 )
@@ -106,7 +118,7 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
                     system_prompt=valid_system_prompt,
                     response_model=str,
                 )
-                followup_args = {"query": query, "answer": answer, "reasoning": reasoning}
+                followup_args = {"query": query, "answer": completion, "reasoning": reasoning}
                 followup_prompt = LLMGateway.render_prompt(
                     filename=self.followup_user_prompt_path, context=followup_args
                 )
@@ -121,4 +133,9 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
                     f"Chain-of-thought: round {round_idx} - follow-up question: {followup_question}"
                 )
 
-        return [answer]
+        if self.save_interaction and context and triplets and completion:
+            await self.save_qa(
+                question=query, answer=completion, context=context_text, triplets=triplets
+            )
+
+        return completion