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]

great_ai/helper/frozen_call.py

@@ -1,11 +1,19 @@
 from copy import deepcopy
 from typing import Any, Dict, Tuple
 
+import numpy as np
 from pydantic import BaseModel
 
 
 def make_hashable(value: Any) -> Any:
     """Build an immutable value key without retaining mutable container contents."""
+    if isinstance(value, np.ndarray):
+        contents = (
+            make_hashable(value.tolist()) if value.dtype.hasobject else value.tobytes()
+        )
+        return type(value), value.dtype, value.shape, contents
+    if isinstance(value, np.generic):
+        return type(value), value.dtype, value.tobytes()
     if isinstance(value, BaseModel):
         return type(value), make_hashable(value.model_dump())
     if isinstance(value, dict):

great_ai/helper/numpy_to_python.py

@@ -0,0 +1,20 @@
+from typing import Any
+
+import numpy as np
+from pydantic_core import PydanticSerializationError
+
+
+def numpy_to_python(value: Any) -> Any:
+    """Convert NumPy values encountered by Pydantic's JSON serializer."""
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+    if isinstance(value, np.floating):
+        return float(value)
+    if isinstance(value, np.generic):
+        converted = value.item()
+        if not isinstance(converted, type(value)):
+            return converted
+    raise PydanticSerializationError(
+        f"Unable to serialize {type(value).__name__}; convert it to a JSON-compatible "
+        "Python value before returning or logging it."
+    )

great_ai/persistence/mongodb_driver.py

@@ -69,7 +69,7 @@ class MongoDbDriver(TracingDatabaseDriver):
         super().configure_credentials()
 
     def save(self, trace: Trace) -> str:
-        serialized = trace.to_flat_dict()
+        serialized = trace.to_flat_dict(mode="json")
         serialized["_id"] = trace.trace_id
 
         with MongoClient[Any](self.mongo_connection_string) as client:
@@ -78,7 +78,7 @@ class MongoDbDriver(TracingDatabaseDriver):
     def save_batch(self, documents: List[Trace]) -> List[str]:
         if not documents:
             return []
-        serialized = [d.to_flat_dict() for d in documents]
+        serialized = [d.to_flat_dict(mode="json") for d in documents]
         for s in serialized:
             s["_id"] = s["trace_id"]
 
@@ -175,7 +175,7 @@ class MongoDbDriver(TracingDatabaseDriver):
         return documents, count
 
     def update(self, id: str, new_version: Trace) -> None:
-        serialized = new_version.to_flat_dict()
+        serialized = new_version.to_flat_dict(mode="json")
         serialized["_id"] = new_version.trace_id
 
         with MongoClient[Any](self.mongo_connection_string) as client:

great_ai/tracing/add_ground_truth.py

@@ -102,7 +102,7 @@ def add_ground_truth(
                 trace_id=str(uuid4()),
                 created=created,
                 original_execution_time_ms=0,
-                logged_values=X if isinstance(X, dict) else {"input": X},
+                logged_values={"input": X},
                 models=[],
                 output=y,
                 feedback=y,