From 58d318dd0e16efa7667b8975ce32fd7fa8979543 Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Sun, 3 Jul 2022 16:17:27 +0200 Subject: [PATCH] Serialize BaseModel --- src/great_ai/great_ai/views/trace.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/great_ai/great_ai/views/trace.py b/src/great_ai/great_ai/views/trace.py index adc744b..287c42a 100644 --- a/src/great_ai/great_ai/views/trace.py +++ b/src/great_ai/great_ai/views/trace.py @@ -2,7 +2,7 @@ from pprint import pformat from typing import Any, Dict, Generic, List, Optional, TypeVar from uuid import uuid4 -from pydantic import Extra, validator +from pydantic import BaseModel, Extra, validator from ..helper import HashableBaseModel from .model import Model @@ -77,7 +77,10 @@ class Trace(Generic[T], HashableBaseModel): "original_execution_time_ms": self.original_execution_time_ms, } ), - **self.logged_values, + **{ + k: v.dict() if isinstance(v, BaseModel) else v + for k, v in self.logged_values.items() + }, "models_flat": self.models_flat, "exception_flat": self.exception_flat, "output_flat": self.output_flat,