Python: resolve string forward references nested inside list/dict/tuple annotations

python/semantic_kernel/schema/kernel_json_schema_builder.py

@@ -1,5 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import ast
 import sys
 import types
 from enum import Enum
@@ -86,6 +87,7 @@ def build_model_schema(
         hints = get_type_hints(model, globalns=model_module_globals, localns={})
 
         for field_name, field_type in hints.items():
+            field_type = cls._resolve_nested_forward_refs(field_type, model_module_globals)
             field_description = None
             if hasattr(model, "model_fields") and field_name in model.model_fields:
                 field_info = model.model_fields[field_name]
@@ -150,6 +152,97 @@ def get_json_schema(cls, parameter_type: type) -> dict[str, Any]:
         type_name = TYPE_MAPPING.get(parameter_type, "object")
         return {"type": type_name}
 
+    _TYPE_EXPRESSION_NODES = (
+        ast.Expression,
+        ast.Name,
+        ast.Attribute,
+        ast.Subscript,
+        ast.Tuple,
+        ast.List,
+        ast.Load,
+        ast.Constant,
+        ast.BinOp,
+        ast.BitOr,
+    )
+
+    @classmethod
+    def _is_type_expression(cls, source: str) -> bool:
+        """Return whether `source` parses as a type expression and nothing more.
+
+        A forward reference names a type: an identifier, a dotted path, a subscription such as
+        `dict[str, Inner]`, or a `X | None` union. Calls, lambdas and comprehensions are not part
+        of that grammar, so rejecting them keeps `eval` from running anything a type annotation
+        would never legitimately contain.
+        """
+        try:
+            tree = ast.parse(source, mode="eval")
+        except SyntaxError:
+            return False
+        for node in ast.walk(tree):
+            if not isinstance(node, cls._TYPE_EXPRESSION_NODES):
+                return False
+            if isinstance(node, ast.Constant) and not isinstance(node.value, (str, int, bool, type(None))):
+                return False
+        return True
+
+    @classmethod
+    def _resolve_nested_forward_refs(cls, annotation: Any, globalns: dict[str, Any]) -> Any:
+        """Resolve string forward references nested inside a generic alias.
+
+        `get_type_hints` evaluates an annotation that *is* a string, but it does not descend into
+        a generic alias that already exists as an object. `list["Inner"]` goes through
+        `list.__class_getitem__`, which stores `"Inner"` verbatim instead of wrapping it in a
+        `ForwardRef`, so nothing resolves it and `build` formats the bare string rather than the
+        class it names.
+
+        Args:
+            annotation: The annotation to resolve, typically a generic alias.
+            globalns: The globals of the module the owning model was defined in.
+
+        Returns:
+            Any: The annotation with resolvable string arguments replaced by the types they name,
+                or the original annotation when nothing could be resolved.
+        """
+        args = get_args(annotation)
+        if not args:
+            return annotation
+
+        resolved_args = []
+        changed = False
+        for arg in args:
+            reference = arg if isinstance(arg, str) else getattr(arg, "__forward_arg__", None)
+            if reference is not None:
+                if not cls._is_type_expression(reference):
+                    # Anything that isn't the grammar of a type expression is not a forward
+                    # reference; refuse to evaluate it rather than widen what a stray
+                    # annotation can run while a schema is being built.
+                    return annotation
+                try:
+                    resolved = eval(reference, globalns, {})
+                except Exception:
+                    # Not resolvable from this module; leave the annotation alone so the existing
+                    # fallback applies instead of raising while a schema is being built.
+                    return annotation
+                changed = True
+            else:
+                resolved = cls._resolve_nested_forward_refs(arg, globalns)
+                changed = changed or resolved is not arg
+            resolved_args.append(resolved)
+
+        if not changed:
+            return annotation
+
+        copy_with = getattr(annotation, "copy_with", None)
+        if copy_with is not None:
+            return copy_with(tuple(resolved_args))
+        origin = get_origin(annotation)
+        if origin is None:
+            return annotation
+        try:
+            return origin[tuple(resolved_args)]
+        except TypeError:
+            return annotation
+
     @classmethod
     def handle_complex_type(
         cls, parameter_type: type, description: str | None = None, structured_output: bool = False

python/tests/unit/schema/test_schema_builder.py

@@ -2,7 +2,7 @@
 
 import json
 from enum import Enum
-from typing import Annotated, Any, Optional, Union
+from typing import Annotated, Any, Optional, Union, get_args
 from unittest.mock import Mock
 
 import pytest
@@ -455,3 +455,113 @@ def test_build_schema_with_nonpydantic_structured_output():
     }
 
     assert structured_output_schema == expected_schema
+
+
+class InnerForward(KernelBaseModel):
+    value: int
+    label: str
+
+
+class HolderForwardList(KernelBaseModel):
+    items: list["InnerForward"] = []
+
+
+class HolderDirectList(KernelBaseModel):
+    items: list[InnerForward] = []
+
+
+class HolderForwardDict(KernelBaseModel):
+    mapping: dict[str, "InnerForward"] = {}
+
+
+class HolderForwardOptional(KernelBaseModel):
+    maybe: Optional["InnerForward"] = None
+
+
+def test_build_list_with_string_forward_reference_matches_direct_reference():
+    """`list["Inner"]` and `list[Inner]` must produce the same schema."""
+    forward = KernelJsonSchemaBuilder.build(HolderForwardList)
+    direct = KernelJsonSchemaBuilder.build(HolderDirectList)
+
+    assert forward == direct
+    assert forward["properties"]["items"]["items"]["properties"] == {
+        "value": {"type": "integer"},
+        "label": {"type": "string"},
+    }
+
+
+def test_build_dict_with_string_forward_reference():
+    schema = KernelJsonSchemaBuilder.build(HolderForwardDict)
+
+    assert schema["properties"]["mapping"]["additionalProperties"]["properties"] == {
+        "value": {"type": "integer"},
+        "label": {"type": "string"},
+    }
+
+
+def test_build_optional_with_string_forward_reference_still_works():
+    """`Optional["Inner"]` already resolved via `get_type_hints`; make sure it still does."""
+    schema = KernelJsonSchemaBuilder.build(HolderForwardOptional)
+
+    assert schema["properties"]["maybe"]["properties"] == {
+        "value": {"type": "integer"},
+        "label": {"type": "string"},
+    }
+
+
+def test_unresolvable_forward_reference_falls_back_instead_of_raising():
+    """An annotation naming something that doesn't exist must not break schema building."""
+
+    class HolderUnknown(KernelBaseModel):
+        items: list["DoesNotExistAnywhere"] = []  # noqa: F821
+
+    schema = KernelJsonSchemaBuilder.build(HolderUnknown)
+
+    assert schema["properties"]["items"]["type"] == "array"
+    assert schema["properties"]["items"]["items"] == {"type": "object"}
+
+
+def test_non_type_expression_forward_reference_is_not_evaluated():
+    """A nested string that isn't a type expression must not be evaluated.
+
+    `list["Inner"]` stores the string verbatim, so whatever it contains reaches the resolver.
+    Only the grammar of a type expression is evaluated; a call is left alone and the annotation
+    comes back unchanged for the existing fallback to handle.
+    """
+    executed = []
+
+    def _canary():
+        executed.append(True)
+        return int
+
+    annotation = list["_canary()"]  # noqa: F821
+    resolved = KernelJsonSchemaBuilder._resolve_nested_forward_refs(annotation, {"_canary": _canary})
+
+    assert executed == []
+    assert resolved is annotation
+
+
+def test_type_expression_forward_reference_is_still_resolved():
+    """The guard must not block the case the resolver exists for."""
+    annotation = list["InnerForward"]  # noqa: F821
+    resolved = KernelJsonSchemaBuilder._resolve_nested_forward_refs(annotation, {"InnerForward": InnerForward})
+
+    assert get_args(resolved) == (InnerForward,)
+
+
+@pytest.mark.parametrize(
+    ("reference", "is_type_expression"),
+    [
+        ("InnerForward", True),
+        ("dict[str, InnerForward]", True),
+        ("InnerForward | None", True),
+        ("tuple[int, str]", True),
+        ("__import__('os').getcwd()", False),
+        ("_canary()", False),
+        ("(lambda: 1)()", False),
+        ("[x for x in ().__class__.__base__.__subclasses__()]", False),
+    ],
+)
+def test_is_type_expression(reference: str, is_type_expression: bool):
+    """The guard admits type expressions and rejects anything that can call out."""
+    assert KernelJsonSchemaBuilder._is_type_expression(reference) is is_type_expression