Use get_type_hints() in function_schema()

Author: Viicos

pydantic_ai_slim/pydantic_ai/_function_schema.py

@@ -7,6 +7,7 @@
 
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
+from functools import partial
 from inspect import Parameter, signature
 from typing import TYPE_CHECKING, Any, Concatenate, cast, get_origin
 
@@ -17,7 +18,7 @@
 from pydantic.json_schema import GenerateJsonSchema
 from pydantic.plugin._schema_validator import create_schema_validator
 from pydantic_core import SchemaValidator, core_schema
-from typing_extensions import ParamSpec, TypeIs, TypeVar
+from typing_extensions import ParamSpec, TypeIs, TypeVar, get_type_hints
 
 from ._griffe import doc_descriptions
 from ._run_context import RunContext
@@ -90,9 +91,6 @@ def function_schema(  # noqa: C901
     Returns:
         A `FunctionSchema` instance.
     """
-    if takes_ctx is None:
-        takes_ctx = _takes_ctx(function)
-
     config = ConfigDict(title=function.__name__, use_attribute_docstrings=True)
     config_wrapper = ConfigWrapper(config)
     gen_schema = _generate_schema.GenerateSchema(config_wrapper)
@@ -103,16 +101,17 @@ def function_schema(  # noqa: C901
     except ValueError as e:
         errors.append(str(e))
         sig = signature(lambda: None)
+    original_func = function.func if isinstance(function, partial) else function
 
-    type_hints = _typing_extra.get_function_type_hints(function)
+    type_hints = get_type_hints(original_func)
 
     var_kwargs_schema: core_schema.CoreSchema | None = None
     fields: dict[str, core_schema.TypedDictField] = {}
     positional_fields: list[str] = []
     var_positional_field: str | None = None
     decorators = _decorators.DecoratorInfos()
 
-    description, field_descriptions = doc_descriptions(function, sig, docstring_format=docstring_format)
+    description, field_descriptions = doc_descriptions(original_func, sig, docstring_format=docstring_format)
 
     if require_parameter_descriptions:
         if takes_ctx:
@@ -127,6 +126,9 @@ def function_schema(  # noqa: C901
             errors.append(f'Missing parameter descriptions for {", ".join(missing_params)}')
 
     for index, (name, p) in enumerate(sig.parameters.items()):
+        if index == 0 and takes_ctx is None:
+            takes_ctx = p.annotation is not sig.empty and _is_call_ctx(type_hints[name])
+
         if p.annotation is sig.empty:
             if takes_ctx and index == 0:
                 # should be the `context` argument, skip
@@ -234,7 +236,7 @@ def function_schema(  # noqa: C901
 TargetCallable = WithCtx[P, R] | WithoutCtx[P, R]
 
 
-def _takes_ctx(callable_obj: TargetCallable[P, R]) -> TypeIs[WithCtx[P, R]]:
+def _takes_ctx(callable_obj: TargetCallable[P, R]) -> TypeIs[WithCtx[P, R]]:  # pyright: ignore[reportUnusedFunction]
     """Check if a callable takes a `RunContext` first argument.
 
     Args: