Fix cognify endpoint (#1105)

<!-- .github/pull_request_template.md -->

## 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.
This commit is contained in:
Igor Ilic 2025-07-18 16:37:04 +02:00 committed by GitHub
parent 875da53c3a
commit 01bab3f0c7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 14 additions and 4 deletions

View file

@ -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

View file

@ -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