diff --git a/mellea/backends/tools.py b/mellea/backends/tools.py index 12b02493c..48890ddd8 100644 --- a/mellea/backends/tools.py +++ b/mellea/backends/tools.py @@ -1344,15 +1344,41 @@ def convert_function_to_ollama_tool( """ doc_string_hash = str(hash(inspect.getdoc(func))) parsed_docstring = _parse_docstring(inspect.getdoc(func)) + # Resolve postponed (string) parameter annotations back to real type + # objects so Pydantic can build schemas for non-builtin parameter types + # under `from __future__ import annotations` (PEP 563). Evaluate + # parameter annotations individually rather than using + # `eval_str=True` because the return annotation is never consumed by + # this schema — resolving it would fail for TYPE_CHECKING-only or + # forward-referenced return types where the pre-existing path + # succeeded. `func.__globals__` already carries `__builtins__`, so no + # defensive copy is needed. + try: + sig = inspect.signature(func, eval_str=True) + except Exception: + sig = inspect.signature(func) + g = getattr(func, "__globals__", {}) + params = [] + for p in sig.parameters.values(): + if isinstance(p.annotation, str): + try: + # ast.literal_eval cannot evaluate type expressions + # (e.g. `Decimal`, `Foo | None`); this mirrors what + # `inspect.signature(..., eval_str=True)` does internally. + p = p.replace(annotation=eval(p.annotation, g)) # noqa: S307 + except Exception: + pass # leave as string; Pydantic will report it + params.append(p) + sig = sig.replace(parameters=params) schema = type( func.__name__, (BaseModel,), { "__annotations__": { k: v.annotation if v.annotation != inspect._empty else str - for k, v in inspect.signature(func).parameters.items() + for k, v in sig.parameters.items() }, - "__signature__": inspect.signature(func), + "__signature__": sig, "__doc__": parsed_docstring[doc_string_hash], }, ).model_json_schema() # type: ignore diff --git a/test/backends/_postponed_annotation_samples.py b/test/backends/_postponed_annotation_samples.py new file mode 100644 index 000000000..431e998af --- /dev/null +++ b/test/backends/_postponed_annotation_samples.py @@ -0,0 +1,66 @@ +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Sample functions with postponed annotations (PEP 563), for regression tests. + +Isolated in its own module because `from __future__ import annotations` is a +module-level directive — it cannot be scoped to a single test function. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from decimal import Decimal + + +@dataclass +class Address: + """A custom, non-builtin parameter type.""" + + city: str + + +@dataclass +class Period: + """A custom parameter type for testing.""" + + name: str + + +def send_letter(to: Address) -> str: + """Send a letter to the given address. + + Args: + to: the destination address + """ + return "sent" + + +def tc_only_return_builtin_param(query: str) -> Decimal: + """TYPE_CHECKING-only return with builtin params. + + Args: + query: query string + """ + return Decimal("0") + + +def tc_return_custom_param(period: Period) -> Decimal: + """TYPE_CHECKING-only return with custom param. + + Args: + period: the period + """ + return Decimal("0") + + +def unresolvable_param(query: NonExistentType) -> str: # type: ignore[name-defined] # noqa: F821 + """Unresolvable parameter annotation. + + Args: + query: query string + """ + return "ok" diff --git a/test/backends/test_discriminated_union_tools.py b/test/backends/test_discriminated_union_tools.py index 88fd28abe..3f3a5a83e 100644 --- a/test/backends/test_discriminated_union_tools.py +++ b/test/backends/test_discriminated_union_tools.py @@ -11,19 +11,34 @@ be preserved and the OAS-3 `discriminator` keyword must be stripped from the output (the JSON Schema subset accepted by tool-calling APIs does not include it; the `Literal` tag fields carry the discriminator signal). + +Also covers `convert_function_to_ollama_tool`'s handling of postponed +annotations (`from __future__ import annotations`, PEP 563): non-builtin +parameter types must resolve to real type objects for Pydantic schema +building, while a return annotation that is unresolvable at call time (e.g. +`TYPE_CHECKING`-only imports) must not break the conversion, since the +produced schema never consumes it. """ import json from typing import Annotated, Literal import pytest -from pydantic import BaseModel, Field, ValidationError +from pydantic import BaseModel, Field, PydanticUserError, ValidationError from mellea.backends.tools import ( MelleaTool, convert_function_to_ollama_tool, validate_tool_arguments, ) +from test.backends._postponed_annotation_samples import ( + Address, + Period, + send_letter, + tc_only_return_builtin_param, + tc_return_custom_param, + unresolvable_param, +) class Cat(BaseModel): @@ -706,5 +721,62 @@ def deeply_nested( ) +class TestPostponedAnnotations: + """Postponed annotation (PEP 563) resolution in generated tool schemas.""" + + def test_resolves_postponed_parameter_annotation(self): + # Regression test: under `from __future__ import annotations`, a + # non-builtin parameter type's annotation is a string rather than the + # real type object, which Pydantic cannot resolve when building the + # dynamic schema model - raising PydanticUserError instead of producing + # a tool schema. + + # Guard the precondition: if the sample module ever drops its + # `from __future__ import annotations`, this test would otherwise keep + # passing without exercising postponed annotations at all. + assert send_letter.__annotations__["to"] == "Address" + + tool = convert_function_to_ollama_tool(send_letter) + assert tool.function is not None + assert tool.function.parameters is not None + + props = tool.function.parameters.model_dump(exclude_none=True)["properties"] + assert props["to"]["type"] == "object" + assert props["to"]["title"] == Address.__name__ + assert props["to"]["properties"]["city"]["type"] == "string" + + def test_type_checking_only_return_with_builtin_params(self): + """TYPE_CHECKING-only return + builtin params must produce schema.""" + tool = convert_function_to_ollama_tool(tc_only_return_builtin_param) + assert tool.function is not None + assert tool.function.parameters is not None + props = tool.function.parameters.model_dump(exclude_none=True)["properties"] + assert "query" in props + + def test_type_checking_only_return_with_custom_param(self): + """TYPE_CHECKING return + custom param must resolve param. + + This is the case that separates the try/except-around-`eval_str=True` + fallback (which discards parameter resolution entirely on any + failure) from per-parameter resolution: the return annotation is + unresolvable, but the custom parameter type still must resolve. + """ + tool = convert_function_to_ollama_tool(tc_return_custom_param) + assert tool.function is not None + assert tool.function.parameters is not None + props = tool.function.parameters.model_dump(exclude_none=True)["properties"] + assert props["period"]["type"] == "object" + assert props["period"]["title"] == Period.__name__ + + def test_unresolvable_parameter_annotation_raises(self): + """Genuinely unresolvable parameter must still raise PydanticUserError. + + Confirms the documented degradation: unresolvable param annotations + surface as `PydanticUserError`, not `NameError`. + """ + with pytest.raises(PydanticUserError): + convert_function_to_ollama_tool(unresolvable_param) + + if __name__ == "__main__": pytest.main([__file__, "-v"])