From 01bab3f0c7dab8747ccbf71d382a82159eaae32a Mon Sep 17 00:00:00 2001 From: Igor Ilic <30923996+dexters1@users.noreply.github.com> Date: Fri, 18 Jul 2025 16:37:04 +0200 Subject: [PATCH] Fix cognify endpoint (#1105) ## Description Have cognify run in background ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin. --- cognee/api/v1/cognify/cognify.py | 15 +++++++++++++-- .../api/v1/cognify/routers/get_cognify_router.py | 3 +-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/cognee/api/v1/cognify/cognify.py b/cognee/api/v1/cognify/cognify.py index 8b9a1be65..7c7821460 100644 --- a/cognee/api/v1/cognify/cognify.py +++ b/cognee/api/v1/cognify/cognify.py @@ -237,8 +237,12 @@ async def run_cognify_as_background_process( graph_db_config: dict = None, vector_db_config: dict = False, ): + # Convert dataset to list if it's a string + if isinstance(datasets, str): + datasets = [datasets] + # Store pipeline status for all pipelines - pipeline_run_started_info = [] + 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 @@ -269,7 +273,14 @@ async def run_cognify_as_background_process( ) # Save dataset Pipeline run started info - pipeline_run_started_info.append(await anext(pipeline_run)) + 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 diff --git a/cognee/api/v1/cognify/routers/get_cognify_router.py b/cognee/api/v1/cognify/routers/get_cognify_router.py index 933b83938..59162382d 100644 --- a/cognee/api/v1/cognify/routers/get_cognify_router.py +++ b/cognee/api/v1/cognify/routers/get_cognify_router.py @@ -32,7 +32,6 @@ logger = get_logger("api.cognify") class CognifyPayloadDTO(InDTO): datasets: Optional[List[str]] = None dataset_ids: Optional[List[UUID]] = None - graph_model: Optional[BaseModel] = KnowledgeGraph run_in_background: Optional[bool] = False @@ -96,7 +95,7 @@ def get_cognify_router() -> APIRouter: datasets = payload.dataset_ids if payload.dataset_ids else payload.datasets cognify_run = await cognee_cognify( - datasets, user, payload.graph_model, run_in_background=payload.run_in_background + datasets, user, run_in_background=payload.run_in_background ) return cognify_run