From bd7a0586b619d484a7102050f9eb81e41c596e23 Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Wed, 28 Oct 2026 22:32:31 +0000 Subject: [PATCH] Bug fixes --- .../routes/bootstrap_prediction_endpoint.py | 89 ++++++++++++++++++- great_ai/helper/frozen_call.py | 8 ++ great_ai/helper/numpy_to_python.py | 20 +++++ great_ai/persistence/mongodb_driver.py | 6 +- great_ai/tracing/add_ground_truth.py | 2 +- 5 files changed, 117 insertions(+), 8 deletions(-) create mode 100644 great_ai/helper/numpy_to_python.py diff --git a/great_ai/deploy/routes/bootstrap_prediction_endpoint.py b/great_ai/deploy/routes/bootstrap_prediction_endpoint.py index f473a2a..266ae48 100644 --- a/great_ai/deploy/routes/bootstrap_prediction_endpoint.py +++ b/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] diff --git a/great_ai/helper/frozen_call.py b/great_ai/helper/frozen_call.py index 2f13e2f..7c9ff9d 100644 --- a/great_ai/helper/frozen_call.py +++ b/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): diff --git a/great_ai/helper/numpy_to_python.py b/great_ai/helper/numpy_to_python.py new file mode 100644 index 0000000..2afbf45 --- /dev/null +++ b/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." + ) diff --git a/great_ai/persistence/mongodb_driver.py b/great_ai/persistence/mongodb_driver.py index 6035e08..27aa0a4 100644 --- a/great_ai/persistence/mongodb_driver.py +++ b/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: diff --git a/great_ai/tracing/add_ground_truth.py b/great_ai/tracing/add_ground_truth.py index 15d87a9..33f68df 100644 --- a/great_ai/tracing/add_ground_truth.py +++ b/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,