Preserve nested model types in prediction requests

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,

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

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",