diff --git a/great_ai/persistence/parallel_tinydb_driver.py b/great_ai/persistence/parallel_tinydb_driver.py index 8a2896f..bf186c5 100644 --- a/great_ai/persistence/parallel_tinydb_driver.py +++ b/great_ai/persistence/parallel_tinydb_driver.py @@ -28,10 +28,10 @@ class ParallelTinyDbDriver(TracingDatabaseDriver): path_to_db = Path(DEFAULT_TRACING_DB_FILENAME) def save(self, trace: Trace) -> str: - return self._safe_execute(lambda db: db.insert(trace.model_dump())) + return self._safe_execute(lambda db: db.insert(trace.model_dump(mode="json"))) def save_batch(self, documents: List[Trace]) -> List[str]: - traces = [d.model_dump() for d in documents] + traces = [d.model_dump(mode="json") for d in documents] return self._safe_execute(lambda db: db.insert_multiple(traces)) def get(self, id: str) -> Optional[Trace]: @@ -100,7 +100,7 @@ class ParallelTinyDbDriver(TracingDatabaseDriver): def update(self, id: str, new_version: Trace) -> None: self._safe_execute( lambda db: db.update( - new_version.model_dump(), lambda d: d["trace_id"] == id + new_version.model_dump(mode="json"), lambda d: d["trace_id"] == id ) ) diff --git a/tests/test_persistence_serialization.py b/tests/test_persistence_serialization.py new file mode 100644 index 0000000..aa24cc0 --- /dev/null +++ b/tests/test_persistence_serialization.py @@ -0,0 +1,36 @@ +from datetime import datetime, timezone +from typing import Set + +from fastapi.testclient import TestClient +from great_ai import GreatAI +from great_ai.context import get_context +from pydantic import BaseModel + + +class DatedInput(BaseModel): + when: datetime + tags: Set[int] + + +def test_tinydb_persists_json_compatible_model_values(isolated_context: None) -> None: + @GreatAI.create + def predict(data: DatedInput) -> int: + return len(data.tags) + + with TestClient(predict.app) as client: + response = client.post( + "/predict", json={"data": {"when": "2026-01-01T00:00:00Z", "tags": [1, 2]}} + ) + assert response.status_code == 200 + trace_id = response.json()["trace_id"] + database = get_context().tracing_database + trace = database.get(trace_id) + assert trace is not None + assert trace.output == 2 + trace.feedback = {"when": datetime(2026, 1, 1, tzinfo=timezone.utc)} + database.update(trace_id, trace) + updated = database.get(trace_id) + assert updated is not None + assert updated.feedback == {"when": "2026-01-01T00:00:00Z"} + database.save_batch([trace.model_copy(update={"trace_id": "batch"})]) + assert database.get("batch") is not None