From ca41feb3285cb85462f4590ffcd69454186338b0 Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Tue, 27 Oct 2026 16:51:58 +0000 Subject: [PATCH] Preserve nested model types in prediction requests --- .../routes/bootstrap_prediction_endpoint.py | 13 +++++----- tests/conftest.py | 23 ++++++++++++++++++ tests/test_prediction_regressions.py | 24 ++++++++++++++++++- 3 files changed, 53 insertions(+), 7 deletions(-) create mode 100644 tests/conftest.py diff --git a/great_ai/deploy/routes/bootstrap_prediction_endpoint.py b/great_ai/deploy/routes/bootstrap_prediction_endpoint.py index b3787ec..f941c62 100644 --- a/great_ai/deploy/routes/bootstrap_prediction_endpoint.py +++ b/great_ai/deploy/routes/bootstrap_prediction_endpoint.py @@ -20,13 +20,14 @@ def bootstrap_prediction_endpoint( @router.post("/predict", status_code=status.HTTP_200_OK, response_model=Trace) async def predict(input_value: schema) -> Trace: # type: ignore try: + # model_dump() recursively converts nested models to dictionaries, + # losing the validated argument types expected by the function. + arguments = { + name: getattr(input_value, name) for name in schema.model_fields + } if inspect.iscoroutinefunction(func): - return await cast(Callable[..., Awaitable[Trace]], func)( - **cast(BaseModel, input_value).model_dump() - ) - return cast(Callable[..., Trace], func)( - **cast(BaseModel, input_value).model_dump() - ) + return await cast(Callable[..., Awaitable[Trace]], func)(**arguments) + return cast(Callable[..., Trace], func)(**arguments) except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..8973cfe --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,23 @@ +from pathlib import Path +from typing import Iterator + +import pytest +from great_ai import RouteConfig, configure, context +from great_ai.large_file import LargeFileBase, LargeFileLocal + + +@pytest.fixture +def isolated_context(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(context, "_context", None) + monkeypatch.setattr(context, "DEFAULT_TRACING_DATABASE_CONFIG_PATHS", {}) + monkeypatch.setattr(context, "DEFAULT_LARGE_FILE_CONFIG_PATHS", {}) + for implementation in (LargeFileBase, LargeFileLocal): + monkeypatch.setattr(implementation, "cache_path", tmp_path / "cache") + monkeypatch.setattr(implementation, "max_cache_size", None) + configure( + disable_se4ml_banner=True, + should_log_exception_stack=False, + route_config=RouteConfig(dashboard_enabled=False), + ) + yield diff --git a/tests/test_prediction_regressions.py b/tests/test_prediction_regressions.py index 3971b11..38e64bd 100644 --- a/tests/test_prediction_regressions.py +++ b/tests/test_prediction_regressions.py @@ -2,7 +2,9 @@ from datetime import datetime from typing import Any, Dict, List from unittest.mock import patch -from great_ai import Trace +import pytest +from fastapi.testclient import TestClient +from great_ai import GreatAI, Trace from great_ai.helper.freeze_arguments import freeze from great_ai.persistence.mongodb_driver import MongoDbDriver from pydantic import BaseModel, Field, field_validator @@ -36,6 +38,26 @@ def test_freezing_preserves_validated_models() -> None: assert type(original) is Payload +@pytest.mark.parametrize("is_async", [False, True]) +def test_http_preserves_nested_models(isolated_context: None, is_async: bool) -> None: + def predict(data: Payload) -> int: + assert isinstance(data.item, Item) + assert all(isinstance(item, Item) for item in data.others) + return data.item.value + sum(item.value for item in data.others) + + async def predict_async(data: Payload) -> int: + return predict(data) + + service = GreatAI.create(predict_async if is_async else predict) + with TestClient(service.app) as client: + response = client.post( + "/predict", + json={"data": {"item": {"number": 1}, "others": [{"number": 2}]}}, + ) + assert response.status_code == 200 + assert response.json()["output"] == 5 + + def test_flattening_preserves_reserved_trace_fields() -> None: trace = Trace( trace_id="actual-id",