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

cognee/modules/pipelines/layers/pipeline_execution_mode.py

@@ -0,0 +1,127 @@
+import asyncio
+from typing import Any, AsyncIterable, AsyncGenerator, Callable, Dict, Union, Awaitable
+from cognee.modules.pipelines.models.PipelineRunInfo import PipelineRunCompleted, PipelineRunErrored
+from cognee.modules.pipelines.queues.pipeline_run_info_queues import push_to_queue
+
+AsyncGenLike = Union[
+    AsyncIterable[Any],
+    AsyncGenerator[Any, None],
+    Callable[..., AsyncIterable[Any]],
+    Callable[..., AsyncGenerator[Any, None]],
+]
+
+
+async def run_pipeline_blocking(pipeline: AsyncGenLike, **params) -> Dict[str, Any]:
+    """
+    Execute a pipeline synchronously (blocking until all results are consumed).
+
+    This function iterates through the given pipeline (an async generator/iterable)
+    until completion, aggregating the run information for each dataset.
+
+    Args:
+        pipeline (AsyncGenLike): The pipeline generator or callable producing async run information.
+        **params: Arbitrary keyword arguments to be passed to the pipeline if it is callable.
+
+    Returns:
+        Dict[str, Any]:
+            - If multiple datasets are processed, a mapping of dataset_id -> last run_info.
+            - If no dataset_id is present in run_info, the run_info itself is returned.
+    """
+    agen = pipeline(**params) if callable(pipeline) else pipeline
+
+    total_run_info: Dict[str, Any] = {}
+
+    async for run_info in agen:
+        dataset_id = getattr(run_info, "dataset_id", None)
+        if dataset_id:
+            total_run_info[dataset_id] = run_info
+        else:
+            total_run_info = run_info
+
+    return total_run_info
+
+
+async def run_pipeline_as_background_process(
+    pipeline: AsyncGenLike,
+    **params,
+) -> Dict[str, Any]:
+    """
+    Execute one or more pipelines as background tasks.
+
+    This function:
+        1. Starts pipelines for each dataset (if multiple datasets are provided).
+        2. Returns the initial "started" run information immediately.
+        3. Continues executing the pipelines in the background,
+           pushing run updates to a queue until each completes.
+
+    Args:
+        pipeline (AsyncGenLike): The pipeline generator or callable producing async run information.
+        **params: Arbitrary keyword arguments to be passed to the pipeline if it is callable.
+                  Expected to include "datasets", which may be a single dataset ID (str)
+                  or a list of dataset IDs.
+
+    Returns:
+        Dict[str, Any]: A mapping of dataset_id -> initial run_info (with payload removed for serialization).
+    """
+
+    datasets = params.get("datasets", None)
+
+    if isinstance(datasets, str):
+        datasets = [datasets]
+
+    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:
+        call_params = dict(params)
+        if "datasets" in call_params:
+            call_params["datasets"] = dataset
+
+        pipeline_run = pipeline(**call_params) if callable(pipeline) else pipeline
+
+        # 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
+
+
+def get_pipeline_executor(
+    run_in_background: bool = False,
+) -> Callable[..., Awaitable[Dict[str, Any]]]:
+    """
+    Return the appropriate pipeline runner.
+
+    Usage:
+        run_fn = get_run_pipeline_fn(run_in_background=True)
+        result = await run_fn(pipeline, **params)
+    """
+    return run_pipeline_as_background_process if run_in_background else run_pipeline_blocking

cognee/modules/pipelines/layers/reset_dataset_pipeline_run_status.py

@@ -0,0 +1,12 @@
+from uuid import UUID
+from cognee.modules.pipelines.methods import get_pipeline_runs_by_dataset, reset_pipeline_run_status
+from cognee.modules.pipelines.models.PipelineRun import PipelineRunStatus
+from cognee.modules.users.models import User
+
+
+async def reset_dataset_pipeline_run_status(dataset_id: UUID, user: User):
+    related_pipeline_runs = await get_pipeline_runs_by_dataset(dataset_id)
+
+    for pipeline_run in related_pipeline_runs:
+        if pipeline_run.status is not PipelineRunStatus.DATASET_PROCESSING_INITIATED:
+            await reset_pipeline_run_status(user.id, dataset_id, pipeline_run.pipeline_name)

cognee/modules/pipelines/layers/resolve_authorized_user_dataset.py

@@ -0,0 +1,34 @@
+from uuid import UUID
+
+from cognee.api.v1.exceptions import DatasetNotFoundError
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_default_user
+from cognee.modules.data.methods import (
+    create_authorized_dataset,
+    get_authorized_dataset,
+    get_authorized_dataset_by_name,
+)
+
+
+async def resolve_authorized_user_dataset(dataset_id: UUID, dataset_name: str, user: User):
+    if not user:
+        user = await get_default_user()
+
+    if dataset_id:
+        authorized_dataset = await get_authorized_dataset(user, dataset_id, "write")
+    elif dataset_name:
+        authorized_dataset = await get_authorized_dataset_by_name(dataset_name, user, "write")
+
+        if not authorized_dataset:
+            authorized_dataset = await create_authorized_dataset(
+                dataset_name=dataset_name, user=user
+            )
+    else:
+        raise ValueError("Either dataset_id or dataset_name must be provided.")
+
+    if not authorized_dataset:
+        raise DatasetNotFoundError(
+            message=f"Dataset ({str(dataset_id) or dataset_name}) not found."
+        )
+
+    return user, authorized_dataset

cognee/modules/pipelines/layers/resolve_authorized_user_datasets.py

@@ -0,0 +1,55 @@
+from uuid import UUID
+from typing import Union, Tuple, List
+
+from cognee.modules.users.methods import get_default_user
+from cognee.modules.users.models import User
+from cognee.modules.data.models import Dataset
+from cognee.modules.data.exceptions import DatasetNotFoundError
+from cognee.modules.data.methods import (
+    get_authorized_existing_datasets,
+    load_or_create_datasets,
+    check_dataset_name,
+)
+
+
+async def resolve_authorized_user_datasets(
+    datasets: Union[str, UUID, list[str], list[UUID]], user: User = None
+) -> Tuple[User, List[Dataset]]:
+    """
+    Function handles creation and dataset authorization if datasets already exist for Cognee.
+    Verifies that provided user has necessary permission for provided Dataset.
+    If Dataset does not exist creates the Dataset and gives permission for the user creating the dataset.
+
+    Args:
+        user: Cognee User request is being processed for, if None default user will be used.
+        datasets: Dataset names or Dataset UUID (in case Datasets already exist)
+
+    Returns:
+
+    """
+    # If no user is provided use default user
+    if user is None:
+        user = await get_default_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.
+        authorized_datasets = existing_datasets
+    else:
+        # If dataset matches an existing Dataset (by name or id), reuse it. Otherwise, create a new Dataset.
+        authorized_datasets = await load_or_create_datasets(datasets, existing_datasets, user)
+
+    if not authorized_datasets:
+        raise DatasetNotFoundError("There are no datasets to work with.")
+
+    for dataset in authorized_datasets:
+        check_dataset_name(dataset.name)
+
+    return user, authorized_datasets

cognee/modules/pipelines/layers/setup_and_check_environment.py

@@ -0,0 +1,41 @@
+import asyncio
+from cognee.context_global_variables import (
+    graph_db_config as context_graph_db_config,
+    vector_db_config as context_vector_db_config,
+)
+
+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,
+)
+
+_first_run_done = False
+_first_run_lock = asyncio.Lock()
+
+
+async def setup_and_check_environment(
+    vector_db_config: dict = None,
+    graph_db_config: dict = None,
+):
+    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()
+
+    global _first_run_done
+    async with _first_run_lock:
+        if not _first_run_done:
+            from cognee.infrastructure.llm.utils import (
+                test_llm_connection,
+                test_embedding_connection,
+            )
+
+            await test_llm_connection()
+            await test_embedding_connection()
+            _first_run_done = True

cognee/modules/pipelines/layers/validate_pipeline_tasks.py

@@ -0,0 +1,20 @@
+from ..tasks.task import Task
+from ..exceptions.tasks import WrongTaskTypeError
+
+
+def validate_pipeline_tasks(tasks: list[Task]):
+    """
+    Validates the tasks argument to ensure it is a list of Task class instances.
+
+    Args:
+        tasks (list[Task]): The list of tasks to be validated.
+    """
+
+    if not isinstance(tasks, list):
+        raise WrongTaskTypeError(f"tasks argument must be a list, got {type(tasks).__name__}.")
+
+    for task in tasks:
+        if not isinstance(task, Task):
+            raise WrongTaskTypeError(
+                f"tasks argument must be a list of Task class instances, got {type(task).__name__} in the list."
+            )

cognee/modules/pipelines/methods/__init__.py

@@ -1,2 +1,4 @@
 from .get_pipeline_run import get_pipeline_run
 from .get_pipeline_run_by_dataset import get_pipeline_run_by_dataset
+from .get_pipeline_runs_by_dataset import get_pipeline_runs_by_dataset
+from .reset_pipeline_run_status import reset_pipeline_run_status

cognee/modules/pipelines/methods/get_pipeline_runs_by_dataset.py

@@ -0,0 +1,34 @@
+from uuid import UUID
+from typing import Optional
+from sqlalchemy import select, func
+from sqlalchemy.orm import aliased
+
+from cognee.infrastructure.databases.relational import get_relational_engine
+from ..models import PipelineRun
+
+
+async def get_pipeline_runs_by_dataset(dataset_id: UUID):
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        query = (
+            select(
+                PipelineRun,
+                func.row_number()
+                .over(
+                    partition_by=(PipelineRun.dataset_id, PipelineRun.pipeline_name),
+                    order_by=PipelineRun.created_at.desc(),
+                )
+                .label("rn"),
+            )
+            .filter(PipelineRun.dataset_id == dataset_id)
+            .subquery()
+        )
+
+        aliased_pipeline_run = aliased(PipelineRun, query)
+
+        latest_run = select(aliased_pipeline_run).filter(query.c.rn == 1)
+
+        runs = (await session.execute(latest_run)).scalars().all()
+
+        return runs

cognee/modules/pipelines/methods/reset_pipeline_run_status.py

@@ -0,0 +1,16 @@
+from uuid import UUID
+from cognee.modules.pipelines.utils.generate_pipeline_id import generate_pipeline_id
+from cognee.modules.pipelines.operations.log_pipeline_run_initiated import (
+    log_pipeline_run_initiated,
+)
+
+
+async def reset_pipeline_run_status(user_id: UUID, dataset_id: UUID, pipeline_name: str):
+    pipeline_id = generate_pipeline_id(user_id, dataset_id, pipeline_name)
+
+    # Without this the pipeline status will be DATASET_PROCESSING_COMPLETED and will skip the execution.
+    await log_pipeline_run_initiated(
+        pipeline_id=pipeline_id,
+        pipeline_name=pipeline_name,
+        dataset_id=dataset_id,
+    )

cognee/modules/pipelines/operations/__init__.py

@@ -2,4 +2,3 @@ from .log_pipeline_run_initiated import log_pipeline_run_initiated
 from .log_pipeline_run_start import log_pipeline_run_start
 from .log_pipeline_run_complete import log_pipeline_run_complete
 from .log_pipeline_run_error import log_pipeline_run_error
-from .pipeline import cognee_pipeline

cognee/modules/pipelines/operations/log_pipeline_run_initiated.py

@@ -4,7 +4,7 @@ from cognee.modules.pipelines.models import PipelineRun, PipelineRunStatus
 from cognee.modules.pipelines.utils import generate_pipeline_run_id
 
 
-async def log_pipeline_run_initiated(pipeline_id: str, pipeline_name: str, dataset_id: UUID):
+async def log_pipeline_run_initiated(pipeline_id: UUID, pipeline_name: str, dataset_id: UUID):
     pipeline_run = PipelineRun(
         pipeline_run_id=generate_pipeline_run_id(pipeline_id, dataset_id),
         pipeline_name=pipeline_name,

cognee/modules/pipelines/operations/pipeline.py

@@ -2,41 +2,22 @@ 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.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.modules.pipelines.layers.resolve_authorized_user_datasets import (
+    resolve_authorized_user_datasets,
 )
-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 +25,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 +35,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
-
-    # If no user is provided use default user
-    if user is None:
-        user = await get_default_user()
+    validate_pipeline_tasks(tasks)
+    await setup_and_check_environment(vector_db_config, graph_db_config)
 
-    # Convert datasets to list
-    if isinstance(datasets, str) or isinstance(datasets, UUID):
-        datasets = [datasets]
+    user, authorized_datasets = await resolve_authorized_user_datasets(datasets, user)
 
-    # 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 +53,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 +62,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/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/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

@@ -29,6 +29,7 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
         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,6 +37,7 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
             top_k=top_k,
             node_type=node_type,
             node_name=node_name,
+            save_interaction=save_interaction,
         )
 
     async def get_completion(
@@ -105,11 +107,16 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
 
             round_idx += 1
 
-        answer = await generate_completion(
+        completion = await generate_completion(
             query=query,
             context=context,
             user_prompt_path=self.user_prompt_path,
             system_prompt_path=self.system_prompt_path,
         )
 
-        return [answer]
+        if self.save_interaction and context and triplets and completion:
+            await self.save_qa(
+                question=query, answer=completion, context=context, triplets=triplets
+            )
+
+        return [completion]