From 5e819525db3cc4b3324050c1f33bf21b96c0d1de Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Tue, 27 Oct 2026 19:53:40 +0000
Subject: [PATCH] Resolve postponed annotations before validation and schema
 creation

---
 .../routes/bootstrap_prediction_endpoint.py   |  4 ++-
 great_ai/helper/get_function_type_hints.py    |  7 +++++
 great_ai/parameters/parameter.py              |  4 +--
 tests/test_argument_regressions.py            | 26 ++++++++++++++++++-
 4 files changed, 37 insertions(+), 4 deletions(-)
 create mode 100644 great_ai/helper/get_function_type_hints.py

diff --git a/great_ai/deploy/routes/bootstrap_prediction_endpoint.py b/great_ai/deploy/routes/bootstrap_prediction_endpoint.py
index f5f09c3..f473a2a 100644
--- a/great_ai/deploy/routes/bootstrap_prediction_endpoint.py
+++ b/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()
diff --git a/great_ai/helper/get_function_type_hints.py b/great_ai/helper/get_function_type_hints.py
new file mode 100644
index 0000000..f6cdae9
--- /dev/null
+++ b/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)
diff --git a/great_ai/parameters/parameter.py b/great_ai/parameters/parameter.py
index df3a7d4..f023be6 100644
--- a/great_ai/parameters/parameter.py
+++ b/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)
diff --git a/tests/test_argument_regressions.py b/tests/test_argument_regressions.py
index d258c23..616db80 100644
--- a/tests/test_argument_regressions.py
+++ b/tests/test_argument_regressions.py
@@ -1,6 +1,8 @@
+import sys
+from types import ModuleType
+
 import pytest
 from fastapi.testclient import TestClient
-
 from great_ai import GreatAI
 
 
@@ -32,3 +34,25 @@ def test_positional_only_inputs_keep_their_bindings(isolated_context: None) -> N
         response = client.post("/predict", json={"first": "b", "second": 3})
         assert response.status_code == 200
         assert response.json()["output"] == "bbb"
+
+
+def test_postponed_model_annotations_resolve_in_their_module(
+    isolated_context: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    module = ModuleType("great_ai_test_annotations")
+    monkeypatch.setitem(sys.modules, module.__name__, module)
+    exec(
+        "from __future__ import annotations\n"
+        "from pydantic import BaseModel\n"
+        "class Payload(BaseModel):\n    value: int\n"
+        "def predict(data: Payload) -> int:\n    return data.value\n",
+        module.__dict__,
+    )
+    service = GreatAI.create(module.predict)
+    with pytest.raises(TypeError):
+        service("wrong")
+    with TestClient(service.app) as client:
+        response = client.post("/predict", json={"data": {"value": 7}})
+        assert response.status_code == 200
+        assert response.json()["output"] == 7
+        assert client.get("/openapi.json").status_code == 200