Feat: log pipeline status and pass it through pipeline [COG-1214]

.pre-commit-config.yaml

@@ -10,7 +10,7 @@ repos:
     -   id: check-added-large-files
 - repo: https://github.com/astral-sh/ruff-pre-commit
   # Ruff version.
-  rev: v0.9.0
+  rev: v0.9.5
   hooks:
     # Run the linter.
     - id: ruff

cognee/api/v1/add/add_v2.py

@@ -9,6 +9,7 @@
 from cognee.infrastructure.databases.vector.pgvector import (
     create_db_and_tables as create_pgvector_db_and_tables,
 )
+from uuid import uuid5, NAMESPACE_OID
 
 
 async def add(
@@ -37,7 +38,8 @@ async def add(
 
     tasks = [Task(resolve_data_directories), Task(ingest_data, dataset_name, user)]
 
-    pipeline = run_tasks(tasks, data, "add_pipeline")
+    dataset_id = uuid5(NAMESPACE_OID, dataset_name)
+    pipeline = run_tasks(tasks=tasks, dataset_id=dataset_id, data=data, pipeline_id="add_pipeline")
 
     async for result in pipeline:
         print(result)

cognee/api/v1/cognify/code_graph_pipeline.py

@@ -69,9 +69,19 @@ async def run_code_graph_pipeline(repo_path, include_docs=True):
             ),
         ]
 
+    pipeline_run_status = None
     if include_docs:
-        async for result in run_tasks(non_code_tasks, repo_path):
-            yield result
+        non_code_pipeline_run = run_tasks(non_code_tasks, None, repo_path, "cognify_pipeline")
+        async for run_status in non_code_pipeline_run:
+            pipeline_run_status = run_status
 
-    async for result in run_tasks(tasks, repo_path, "cognify_code_pipeline"):
-        yield result
+    from cognee.modules.data.methods import get_datasets
+
+    existing_datasets = await get_datasets(user.id)
+    code_pipeline_run = run_tasks(
+        tasks, existing_datasets[0].id, repo_path, "cognify_code_pipeline"
+    )
+    async for run_status in code_pipeline_run:
+        pipeline_run_status = run_status
+
+    return pipeline_run_status

cognee/api/v1/cognify/cognify_v2.py

@@ -12,7 +12,6 @@
 from cognee.modules.pipelines import run_tasks
 from cognee.modules.pipelines.models import PipelineRunStatus
 from cognee.modules.pipelines.operations.get_pipeline_status import get_pipeline_status
-from cognee.modules.pipelines.operations.log_pipeline_status import log_pipeline_status
 from cognee.modules.pipelines.tasks.Task import Task
 from cognee.modules.users.methods import get_default_user
 from cognee.modules.users.models import User
@@ -73,8 +72,6 @@ async def cognify(
 async def run_cognify_pipeline(dataset: Dataset, user: User, tasks: list[Task]):
     data_documents: list[Data] = await get_dataset_data(dataset_id=dataset.id)
 
-    document_ids_str = [str(document.id) for document in data_documents]
-
     dataset_id = dataset.id
     dataset_name = generate_dataset_name(dataset.name)
 
@@ -84,21 +81,12 @@ async def run_cognify_pipeline(dataset: Dataset, user: User, tasks: list[Task]):
     task_status = await get_pipeline_status([dataset_id])
 
     if (
-        dataset_id in task_status
-        and task_status[dataset_id] == PipelineRunStatus.DATASET_PROCESSING_STARTED
+        str(dataset_id) in task_status
+        and task_status[str(dataset_id)] == PipelineRunStatus.DATASET_PROCESSING_STARTED
     ):
         logger.info("Dataset %s is already being processed.", dataset_name)
         return
 
-    await log_pipeline_status(
-        dataset_id,
-        PipelineRunStatus.DATASET_PROCESSING_STARTED,
-        {
-            "dataset_name": dataset_name,
-            "files": document_ids_str,
-        },
-    )
-
     try:
         if not isinstance(tasks, list):
             raise ValueError("Tasks must be a list")
@@ -107,34 +95,19 @@ async def run_cognify_pipeline(dataset: Dataset, user: User, tasks: list[Task]):
             if not isinstance(task, Task):
                 raise ValueError(f"Task {task} is not an instance of Task")
 
-        pipeline = run_tasks(tasks, data_documents, "cognify_pipeline")
+        pipeline_run = run_tasks(tasks, dataset.id, data_documents, "cognify_pipeline")
+        pipeline_run_status = None
 
-        async for result in pipeline:
-            print(result)
+        async for run_status in pipeline_run:
+            pipeline_run_status = run_status
 
         await index_graph_edges()
 
         send_telemetry("cognee.cognify EXECUTION COMPLETED", user.id)
+        return pipeline_run_status
 
-        await log_pipeline_status(
-            dataset_id,
-            PipelineRunStatus.DATASET_PROCESSING_COMPLETED,
-            {
-                "dataset_name": dataset_name,
-                "files": document_ids_str,
-            },
-        )
     except Exception as error:
         send_telemetry("cognee.cognify EXECUTION ERRORED", user.id)
-
-        await log_pipeline_status(
-            dataset_id,
-            PipelineRunStatus.DATASET_PROCESSING_ERRORED,
-            {
-                "dataset_name": dataset_name,
-                "files": document_ids_str,
-            },
-        )
         raise error
 
 

cognee/modules/pipelines/models/PipelineRun.py

@@ -1,7 +1,7 @@
 import enum
 from uuid import uuid4
 from datetime import datetime, timezone
-from sqlalchemy import Column, DateTime, JSON, Enum, UUID
+from sqlalchemy import Column, DateTime, JSON, Enum, UUID, String
 from cognee.infrastructure.databases.relational import Base
 
 
@@ -20,5 +20,5 @@ class PipelineRun(Base):
 
     status = Column(Enum(PipelineRunStatus))
 
-    run_id = Column(UUID, index=True)
+    pipeline_id = Column(String, index=True)
     run_info = Column(JSON)

cognee/modules/pipelines/operations/__init__.py

@@ -0,0 +1,3 @@
+from .logPipelineRunStart import logPipelineRunStart
+from .logPipelineRunComplete import logPipelineRunComplete
+from .logPipelineRunError import logPipelineRunError

cognee/modules/pipelines/operations/get_pipeline_status.py

@@ -7,38 +7,33 @@
 
 async def get_pipeline_status(pipeline_ids: list[UUID]):
     db_engine = get_relational_engine()
+    dialect = db_engine.engine.dialect.name
 
     async with db_engine.get_async_session() as session:
+        if dialect == "sqlite":
+            dataset_id_column = func.json_extract(PipelineRun.run_info, "$.dataset_id")
+        else:
+            dataset_id_column = PipelineRun.run_info.op("->>")("dataset_id")
+
         query = (
             select(
                 PipelineRun,
                 func.row_number()
                 .over(
-                    partition_by=PipelineRun.run_id,
+                    partition_by=dataset_id_column,
                     order_by=PipelineRun.created_at.desc(),
                 )
                 .label("rn"),
             )
-            .filter(PipelineRun.run_id.in_(pipeline_ids))
+            .filter(dataset_id_column.in_([str(id) for id in pipeline_ids]))
             .subquery()
         )
 
         aliased_pipeline_run = aliased(PipelineRun, query)
-
         latest_runs = select(aliased_pipeline_run).filter(query.c.rn == 1)
 
         runs = (await session.execute(latest_runs)).scalars().all()
 
-        pipeline_statuses = {str(run.run_id): run.status for run in runs}
+        pipeline_statuses = {run.run_info["dataset_id"]: run.status for run in runs}
 
         return pipeline_statuses
-
-        # f"""SELECT data_id, status
-        # FROM (
-        #     SELECT data_id, status, ROW_NUMBER() OVER (PARTITION BY data_id ORDER BY created_at DESC) as rn
-        #     FROM cognee.cognee.task_runs
-        #     WHERE data_id IN ({formatted_data_ids})
-        # ) t
-        # WHERE rn = 1;"""
-
-    # return { dataset["data_id"]: dataset["status"] for dataset in datasets_statuses }

cognee/modules/pipelines/operations/logPipelineRunComplete.py

@@ -0,0 +1,34 @@
+from uuid import UUID, uuid4
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.data.models import Data
+from cognee.modules.pipelines.models import PipelineRun, PipelineRunStatus
+from typing import Any
+
+
+async def logPipelineRunComplete(pipeline_id: UUID, dataset_id: UUID, data: Any):
+    if not data:
+        data_info = "None"
+    elif isinstance(data, list) and all(isinstance(item, Data) for item in data):
+        data_info = [str(item.id) for item in data]
+    else:
+        data_info = data
+
+    pipeline_run_id = uuid4()
+
+    pipeline_run = PipelineRun(
+        id=pipeline_run_id,
+        pipeline_id=pipeline_id,
+        status=PipelineRunStatus.DATASET_PROCESSING_COMPLETED,
+        run_info={
+            "dataset_id": str(dataset_id),
+            "data": data_info,
+        },
+    )
+
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        session.add(pipeline_run)
+        await session.commit()
+
+    return pipeline_run

cognee/modules/pipelines/operations/logPipelineRunError.py

@@ -0,0 +1,35 @@
+from uuid import UUID, uuid4
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.data.models import Data
+from cognee.modules.pipelines.models import PipelineRun, PipelineRunStatus
+from typing import Any
+
+
+async def logPipelineRunError(pipeline_id: str, dataset_id: UUID, data: Any, e: Exception):
+    if not data:
+        data_info = "None"
+    elif isinstance(data, list) and all(isinstance(item, Data) for item in data):
+        data_info = [str(item.id) for item in data]
+    else:
+        data_info = data
+
+    pipeline_run_id = uuid4()
+
+    pipeline_run = PipelineRun(
+        id=pipeline_run_id,
+        pipeline_id=pipeline_id,
+        status=PipelineRunStatus.DATASET_PROCESSING_ERRORED,
+        run_info={
+            "dataset_id": str(dataset_id),
+            "data": data_info,
+            "error": str(e),
+        },
+    )
+
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        session.add(pipeline_run)
+        await session.commit()
+
+    return pipeline_run

cognee/modules/pipelines/operations/logPipelineRunStart.py

@@ -0,0 +1,34 @@
+from uuid import UUID, uuid4
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.data.models import Data
+from cognee.modules.pipelines.models import PipelineRun, PipelineRunStatus
+from typing import Any
+
+
+async def logPipelineRunStart(pipeline_id: str, dataset_id: UUID, data: Any):
+    if not data:
+        data_info = "None"
+    elif isinstance(data, list) and all(isinstance(item, Data) for item in data):
+        data_info = [str(item.id) for item in data]
+    else:
+        data_info = data
+
+    pipeline_run_id = uuid4()
+
+    pipeline_run = PipelineRun(
+        id=pipeline_run_id,
+        pipeline_id=pipeline_id,
+        status=PipelineRunStatus.DATASET_PROCESSING_STARTED,
+        run_info={
+            "dataset_id": str(dataset_id),
+            "data": data_info,
+        },
+    )
+
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        session.add(pipeline_run)
+        await session.commit()
+
+    return pipeline_run

cognee/modules/pipelines/operations/run_tasks.py

@@ -1,7 +1,14 @@
 import inspect
 import json
 import logging
-
+from uuid import UUID
+
+from typing import Any
+from cognee.modules.pipelines.operations import (
+    logPipelineRunStart,
+    logPipelineRunComplete,
+    logPipelineRunError,
+)
 from cognee.modules.settings import get_current_settings
 from cognee.modules.users.methods import get_default_user
 from cognee.modules.users.models import User
@@ -261,6 +268,17 @@ async def run_tasks_with_telemetry(tasks: list[Task], data, pipeline_name: str):
         raise error
 
 
-async def run_tasks(tasks: list[Task], data=None, pipeline_name: str = "default_pipeline"):
-    async for result in run_tasks_with_telemetry(tasks, data, pipeline_name):
-        yield result
+async def run_tasks(tasks: list[Task], dataset_id: UUID, data: Any, pipeline_id: str):
+    pipeline_run = await logPipelineRunStart(pipeline_id, dataset_id, data)
+
+    yield pipeline_run
+
+    try:
+        async for _ in run_tasks_with_telemetry(tasks, data, pipeline_id):
+            pass
+
+        yield await logPipelineRunComplete(pipeline_id, dataset_id, data)
+
+    except Exception as e:
+        yield await logPipelineRunError(pipeline_id, dataset_id, data, e)
+        raise e

cognee/tests/integration/run_toy_tasks/run_task_from_queue_test.py

@@ -1,8 +1,9 @@
 import asyncio
 from queue import Queue
 
-from cognee.modules.pipelines.operations.run_tasks import run_tasks
+from cognee.modules.pipelines.operations.run_tasks import run_tasks_base
 from cognee.modules.pipelines.tasks.Task import Task
+from cognee.modules.users.methods import get_default_user
 
 
 async def pipeline(data_queue):
@@ -19,13 +20,15 @@ async def add_one(num):
     async def multiply_by_two(num):
         yield num * 2
 
-    tasks_run = run_tasks(
+    user = await get_default_user()
+    tasks_run = run_tasks_base(
         [
             Task(queue_consumer),
             Task(add_one),
             Task(multiply_by_two),
         ],
-        pipeline_name="test_run_tasks_from_queue",
+        data=None,
+        user=user,
     )
 
     results = [2, 4, 6, 8, 10, 12, 14, 16, 18, 20]
@@ -50,3 +53,7 @@ async def queue_producer():
 
 def test_run_tasks_from_queue():
     asyncio.run(run_queue())
+
+
+if __name__ == "__main__":
+    asyncio.run(run_queue())