Resolve postponed annotations before validation and schema creation

great_ai/deploy/routes/bootstrap_prediction_endpoint.py

@@ -5,6 +5,7 @@ from fastapi import APIRouter, FastAPI, HTTPException, status
 from pydantic import BaseModel, create_model
 
 from ...helper import get_function_metadata_store
+from ...helper.get_function_type_hints import get_function_type_hints
 from ...views import Trace
 
 
@@ -47,9 +48,10 @@ def bootstrap_prediction_endpoint(
 
 def _get_schema(func: Callable) -> Type[BaseModel]:
     signature = inspect.signature(func)
+    type_hints = get_function_type_hints(func)
     parameters = {
         p.name: (
-            p.annotation if p.annotation != inspect._empty else Any,
+            type_hints.get(p.name, Any),
             p.default if p.default != inspect._empty else ...,
         )
         for p in signature.parameters.values()

great_ai/helper/get_function_type_hints.py

@@ -0,0 +1,7 @@
+import inspect
+from typing import Any, Callable, Dict, get_type_hints
+
+
+def get_function_type_hints(func: Callable) -> Dict[str, Any]:
+    """Resolve postponed annotations in the original function's namespace."""
+    return get_type_hints(inspect.unwrap(func), include_extras=True)

great_ai/parameters/parameter.py

@@ -6,6 +6,7 @@ from typeguard import TypeCheckError, check_type
 from ..errors import ArgumentValidationError
 from ..helper import get_arguments, get_function_metadata_store
 from ..helper.assert_function_is_not_finalised import assert_function_is_not_finalised
+from ..helper.get_function_type_hints import get_function_type_hints
 from ..tracing.tracing_context import TracingContext
 
 F = TypeVar("F", bound=Callable)
@@ -56,14 +57,13 @@ def parameter(
         assert_function_is_not_finalised(func)
 
         actual_name = f"arg:{parameter_name}"
+        expected_type = get_function_type_hints(func).get(parameter_name)
 
         @wraps(func)
         def wrapper(*args: Any, **kwargs: Dict[str, Any]) -> Any:
             arguments = get_arguments(func, args, kwargs)
             argument = arguments.get(parameter_name)
 
-            expected_type = func.__annotations__.get(parameter_name)
-
             if expected_type is not None:
                 try:
                     check_type(argument, expected_type)