Bug fixes

great_ai/deploy/routes/bootstrap_prediction_endpoint.py

@@ -1,11 +1,33 @@
 import inspect
-from typing import Any, Awaitable, Callable, Type, Union, cast
+from types import UnionType
+from typing import (
+    Annotated,
+    Any,
+    Awaitable,
+    Callable,
+    Type,
+    Union,
+    cast,
+    get_args,
+    get_origin,
+)
 
+import numpy as np
 from fastapi import APIRouter, FastAPI, HTTPException, status
-from pydantic import BaseModel, create_model
+from numpy.typing import NDArray
+from pydantic import (
+    AfterValidator,
+    BaseModel,
+    PlainSerializer,
+    PlainValidator,
+    WithJsonSchema,
+    create_model,
+)
 
+from ...errors import ArgumentValidationError
 from ...helper import get_function_metadata_store
 from ...helper.get_function_type_hints import get_function_type_hints
+from ...helper.numpy_to_python import numpy_to_python
 from ...views import Trace
 
 
@@ -17,13 +39,17 @@ def bootstrap_prediction_endpoint(
     )
 
     schema = _get_schema(func)
+    return_type = _get_api_type(get_function_type_hints(func).get("return", Any))
+    response_model = Trace[return_type]  # type: ignore[valid-type]
     positional_names = [
         p.name
         for p in inspect.signature(func).parameters.values()
         if p.kind == p.POSITIONAL_ONLY and p.name in schema.model_fields
     ]
 
-    @router.post("/predict", status_code=status.HTTP_200_OK, response_model=Trace)
+    @router.post(
+        "/predict", status_code=status.HTTP_200_OK, response_model=response_model
+    )
     async def predict(input_value: schema) -> Trace:  # type: ignore
         try:
             # model_dump() recursively converts nested models to dictionaries,
@@ -37,6 +63,11 @@ def bootstrap_prediction_endpoint(
                     *positional, **arguments
                 )
             return cast(Callable[..., Trace], func)(*positional, **arguments)
+        except ArgumentValidationError as e:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
+                detail=str(e),
+            ) from e
         except Exception as e:
             raise HTTPException(
                 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
@@ -51,7 +82,7 @@ def _get_schema(func: Callable) -> Type[BaseModel]:
     type_hints = get_function_type_hints(func)
     parameters = {
         p.name: (
-            type_hints.get(p.name, Any),
+            _get_api_type(type_hints.get(p.name, Any)),
             p.default if p.default != inspect._empty else ...,
         )
         for p in signature.parameters.values()
@@ -60,3 +91,53 @@ def _get_schema(func: Callable) -> Type[BaseModel]:
 
     schema: Type[BaseModel] = create_model("InputModel", **parameters)  # type: ignore
     return schema
+
+
+def _validate_numpy_array(value: Any) -> np.ndarray:
+    if not isinstance(value, (list, np.ndarray)):
+        raise ValueError("Expected an array (a JSON list or NumPy ndarray)")
+    return np.asarray(value)
+
+
+_NumpyArray = Annotated[
+    np.ndarray,
+    PlainValidator(_validate_numpy_array),
+    PlainSerializer(lambda value: value.tolist(), when_used="json"),
+    WithJsonSchema({"type": "array", "items": {}}),
+]
+
+
+def _get_api_type(annotation: Any) -> Any:
+    """Represent NumPy arrays as JSON arrays without losing other annotations."""
+    origin = get_origin(annotation)
+    if annotation in (np.ndarray, NDArray) or origin in (np.ndarray, NDArray):
+        return _NumpyArray
+    if isinstance(annotation, type) and issubclass(annotation, np.generic):
+        scalar_types = (
+            (np.bool_, bool),
+            (np.integer, int),
+            (np.floating, float),
+            (np.complexfloating, complex),
+            (np.str_, str),
+            (np.bytes_, bytes),
+        )
+        for numpy_type, python_type in scalar_types:
+            if issubclass(annotation, numpy_type):
+                return Annotated[
+                    python_type,
+                    AfterValidator(annotation),
+                    PlainSerializer(
+                        numpy_to_python, return_type=python_type, when_used="json"
+                    ),
+                ]
+    arguments = get_args(annotation)
+    if origin is Annotated:
+        return Annotated[_get_api_type(arguments[0]), *arguments[1:]]
+    converted = tuple(_get_api_type(argument) for argument in arguments)
+    if all(old is new for old, new in zip(arguments, converted)):
+        return annotation
+    if origin in (Union, UnionType):
+        return Union[converted]
+    if hasattr(annotation, "copy_with"):
+        return annotation.copy_with(converted)
+    return origin[converted]