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

cognee/modules/notebooks/operations/run_in_local_sandbox.py

@@ -0,0 +1,55 @@
+import io
+import sys
+import traceback
+
+
+def wrap_in_async_handler(user_code: str) -> str:
+    return (
+        "from cognee.infrastructure.utils.run_sync import run_sync\n\n"
+        "async def __user_main__():\n"
+        + "\n".join("    " + line for line in user_code.strip().split("\n"))
+        + "\n"
+        "    globals().update(locals())\n\n"
+        "run_sync(__user_main__())\n"
+    )
+
+
+def run_in_local_sandbox(code, environment=None):
+    environment = environment or {}
+    code = wrap_in_async_handler(code.replace("\xa0", "\n"))
+
+    buffer = io.StringIO()
+    sys_stdout = sys.stdout
+    sys.stdout = buffer
+    sys.stderr = buffer
+
+    error = None
+
+    printOutput = []
+
+    def customPrintFunction(output):
+        printOutput.append(output)
+
+    environment["print"] = customPrintFunction
+
+    try:
+        exec(code, environment)
+    except Exception:
+        error = traceback.format_exc()
+    finally:
+        sys.stdout = sys_stdout
+        sys.stderr = sys_stdout
+
+    return printOutput, error
+
+
+if __name__ == "__main__":
+    run_in_local_sandbox("""
+import cognee
+
+await cognee.add("Test file with some random content 3.")
+
+a = "asd"
+
+b = {"c": "dfgh"}
+""")

cognee/modules/pipelines/__init__.py

@@ -1,4 +1,4 @@
 from .tasks.task import Task
 from .operations.run_tasks import run_tasks
 from .operations.run_parallel import run_tasks_parallel
-from .operations.pipeline import cognee_pipeline
+from .operations.pipeline import run_pipeline

cognee/modules/pipelines/exceptions/tasks.py

@@ -0,0 +1,18 @@
+from fastapi import status
+from cognee.exceptions import CogneeValidationError
+
+
+class WrongTaskTypeError(CogneeValidationError):
+    """
+    Raised when the tasks argument is not a list of Task class instances.
+    """
+
+    def __init__(
+        self,
+        message: str = "tasks argument must be a list, containing Task class instances.",
+        name: str = "WrongTaskTypeError",
+        status_code=status.HTTP_400_BAD_REQUEST,
+    ):
+        self.message = message
+        self.name = name
+        self.status_code = status_code

cognee/modules/pipelines/layers/__init__.py

@@ -0,0 +1 @@
+from .validate_pipeline_tasks import validate_pipeline_tasks

cognee/modules/pipelines/layers/check_pipeline_run_qualification.py

@@ -0,0 +1,59 @@
+from typing import Union, Optional
+from cognee.modules.data.models import Dataset
+from cognee.modules.data.models import Data
+from cognee.modules.pipelines.models import PipelineRunStatus
+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.shared.logging_utils import get_logger
+
+from cognee.modules.pipelines.models.PipelineRunInfo import (
+    PipelineRunCompleted,
+    PipelineRunStarted,
+)
+
+logger = get_logger(__name__)
+
+
+async def check_pipeline_run_qualification(
+    dataset: Dataset, data: list[Data], pipeline_name: str
+) -> Optional[Union[PipelineRunStarted, PipelineRunCompleted]]:
+    """
+    Function used to determine if pipeline is currently being processed or was already processed.
+    In case pipeline was or is being processed return value is returned and current pipline execution should be stopped.
+    In case pipeline is not or was not processed there will be no return value and pipeline processing can start.
+
+    Args:
+        dataset: Dataset object
+        data: List of Data
+        pipeline_name: pipeline name
+
+    Returns: Pipeline state if it is being processed or was already processed
+
+    """
+
+    # 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 = {}
+
+    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)
+            return PipelineRunStarted(
+                pipeline_run_id=pipeline_run.pipeline_run_id,
+                dataset_id=dataset.id,
+                dataset_name=dataset.name,
+                payload=data,
+            )
+        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)
+            return PipelineRunCompleted(
+                pipeline_run_id=pipeline_run.pipeline_run_id,
+                dataset_id=dataset.id,
+                dataset_name=dataset.name,
+            )
+
+    return

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,28 @@
+from uuid import UUID
+from typing import Optional, List
+
+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, pipeline_names: Optional[list[str]] = None
+):
+    """Reset the status of all (or selected) pipeline runs for a dataset.
+
+    If *pipeline_names* is given, only runs whose *pipeline_name* is in
+    that list are touched.
+    """
+    related_pipeline_runs = await get_pipeline_runs_by_dataset(dataset_id)
+
+    for pipeline_run in related_pipeline_runs:
+        # Skip runs that are initiated
+        if pipeline_run.status is PipelineRunStatus.DATASET_PROCESSING_INITIATED:
+            continue
+
+        # If a name filter is provided, skip non-matching runs
+        if pipeline_names is not None and pipeline_run.pipeline_name not in pipeline_names:
+            continue
+
+        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,