From a9668a6f37bdec74cda6844343755c073a89d123 Mon Sep 17 00:00:00 2001 From: estelle Date: Fri, 23 May 2025 17:02:47 +0200 Subject: [PATCH] Fix type so that subclasses are dumped with all their fields (was causing issues in logging with prettify function that calls model_dump where all fields were removed because not in DataModel) --- .../experimental/pipeline/types/orchestration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/neo4j_graphrag/experimental/pipeline/types/orchestration.py b/src/neo4j_graphrag/experimental/pipeline/types/orchestration.py index 1e1497ae3..ba6f99636 100644 --- a/src/neo4j_graphrag/experimental/pipeline/types/orchestration.py +++ b/src/neo4j_graphrag/experimental/pipeline/types/orchestration.py @@ -18,7 +18,7 @@ import enum from typing import Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, SerializeAsAny from neo4j_graphrag.experimental.pipeline.component import DataModel @@ -40,7 +40,7 @@ def possible_next_status(self) -> list[RunStatus]: class RunResult(BaseModel): status: RunStatus = RunStatus.DONE - result: Optional[DataModel] = None + result: Optional[SerializeAsAny[DataModel]] = None timestamp: datetime.datetime = Field( default_factory=lambda: datetime.datetime.now(datetime.timezone.utc) )