Add resolver dependency injection for MCPServer tools (#2969)
Co-authored-by: Max Isbey <224885523+maxisbey@users.noreply.github.com>
This commit is contained in:
parent
4b519782f1
commit
f664db8952
6 changed files with 1844 additions and 15 deletions
|
|
@ -11,6 +11,7 @@ from mcp_types.v2025_11_25 import PrimitiveSchemaDefinition
|
|||
from pydantic import BaseModel, ValidationError
|
||||
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
|
||||
from pydantic_core import core_schema
|
||||
from typing_extensions import TypeAliasType
|
||||
|
||||
from mcp.server.session import ServerSession
|
||||
|
||||
|
|
@ -36,7 +37,11 @@ class CancelledElicitation(BaseModel):
|
|||
action: Literal["cancel"] = "cancel"
|
||||
|
||||
|
||||
ElicitationResult = AcceptedElicitation[ElicitSchemaModelT] | DeclinedElicitation | CancelledElicitation
|
||||
ElicitationResult = TypeAliasType(
|
||||
"ElicitationResult",
|
||||
AcceptedElicitation[ElicitSchemaModelT] | DeclinedElicitation | CancelledElicitation,
|
||||
type_params=(ElicitSchemaModelT,),
|
||||
)
|
||||
|
||||
|
||||
class AcceptedUrlElicitation(BaseModel):
|
||||
|
|
|
|||
|
|
@ -5,6 +5,14 @@ from mcp_types import Icon
|
|||
from mcp.server.extension import Extension, MethodBinding, ResourceBinding, ToolBinding
|
||||
|
||||
from .context import Context
|
||||
from .resolve import (
|
||||
AcceptedElicitation,
|
||||
CancelledElicitation,
|
||||
DeclinedElicitation,
|
||||
Elicit,
|
||||
ElicitationResult,
|
||||
Resolve,
|
||||
)
|
||||
from .resources import DEFAULT_RESOURCE_SECURITY, ResourceSecurity
|
||||
from .server import MCPServer, require_client_extension
|
||||
from .utilities.types import Audio, Image
|
||||
|
|
@ -15,6 +23,12 @@ __all__ = [
|
|||
"Image",
|
||||
"Audio",
|
||||
"Icon",
|
||||
"Resolve",
|
||||
"Elicit",
|
||||
"ElicitationResult",
|
||||
"AcceptedElicitation",
|
||||
"DeclinedElicitation",
|
||||
"CancelledElicitation",
|
||||
"Extension",
|
||||
"ToolBinding",
|
||||
"ResourceBinding",
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
from typing import TYPE_CHECKING, Any, Generic
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Generic, cast
|
||||
|
||||
from mcp_types import ClientCapabilities, InputResponseRequestParams, InputResponses, LoggingLevel
|
||||
from pydantic import AnyUrl, BaseModel
|
||||
|
|
@ -217,6 +217,16 @@ class Context(BaseModel, Generic[LifespanContextT, RequestT]):
|
|||
"""
|
||||
return self.request_context.meta.get("client_id") if self.request_context.meta else None # pragma: no cover
|
||||
|
||||
@property
|
||||
def headers(self) -> Mapping[str, str] | None:
|
||||
"""Request headers carried by this message, when the transport has them.
|
||||
|
||||
Populated by HTTP-based transports; `None` on stdio or when the
|
||||
transport's request object carries no headers. Headers are
|
||||
client-supplied input - never treat one as an identity assertion.
|
||||
"""
|
||||
return cast("Mapping[str, str] | None", getattr(self.request_context.request, "headers", None))
|
||||
|
||||
@property
|
||||
def request_id(self) -> str:
|
||||
"""Get the unique ID for this request."""
|
||||
|
|
|
|||
324
src/mcp/server/mcpserver/resolve.py
Normal file
324
src/mcp/server/mcpserver/resolve.py
Normal file
|
|
@ -0,0 +1,324 @@
|
|||
"""Resolver dependency injection for MCPServer tools.
|
||||
|
||||
A tool parameter annotated `Annotated[T, Resolve(fn)]` is filled by running the
|
||||
resolver `fn` before the tool body, instead of from the LLM-supplied arguments.
|
||||
Resolvers form a DAG: a resolver may declare its own `Resolve(...)` dependencies,
|
||||
take tool arguments by name, and take the `Context`. A resolver may return
|
||||
`Elicit[T]` to ask the client; the framework runs the elicitation and injects the
|
||||
answer.
|
||||
|
||||
Whether the consumer receives the unwrapped model or the full
|
||||
`ElicitationResult` union is decided by the consumer's annotation:
|
||||
|
||||
- `Annotated[T, Resolve(fn)]` -> unwrapped `T`; decline/cancel aborts the call.
|
||||
- `Annotated[ElicitationResult[T], Resolve(fn)]` (or a specific member) -> the
|
||||
full outcome; the consumer branches on accept/decline/cancel.
|
||||
|
||||
Each resolver runs at most once per `tools/call` (memoized by function identity).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import typing
|
||||
from collections.abc import Callable, Hashable, Mapping
|
||||
from typing import Annotated, Any, Generic, cast, get_args, get_origin
|
||||
|
||||
import anyio.to_thread
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
from mcp.server.elicitation import (
|
||||
AcceptedElicitation,
|
||||
CancelledElicitation,
|
||||
DeclinedElicitation,
|
||||
ElicitationResult,
|
||||
)
|
||||
from mcp.server.mcpserver.context import Context
|
||||
from mcp.server.mcpserver.exceptions import InvalidSignature, ToolError
|
||||
from mcp.shared._callable_inspection import is_async_callable
|
||||
|
||||
T = TypeVar("T", bound=BaseModel)
|
||||
|
||||
# The union members the framework injects when a consumer opts into the outcome.
|
||||
_ELICITATION_RESULT_MEMBERS = (AcceptedElicitation, DeclinedElicitation, CancelledElicitation)
|
||||
|
||||
|
||||
class Resolve:
|
||||
"""Marker for `Annotated[T, Resolve(fn)]`: fill the parameter by running `fn`."""
|
||||
|
||||
def __init__(self, fn: Callable[..., Any]) -> None:
|
||||
self.fn = fn
|
||||
|
||||
|
||||
class Elicit(Generic[T]):
|
||||
"""A resolver's request to ask the client.
|
||||
|
||||
Returned from a resolver to signal that the value must be elicited. The
|
||||
framework runs `ctx.elicit(message, schema)` and injects the outcome.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str, schema: type[T]) -> None:
|
||||
self.message = message
|
||||
self.schema = schema
|
||||
|
||||
|
||||
class _ParamPlan:
|
||||
"""How to fill one resolver parameter, decided once at registration."""
|
||||
|
||||
kind: str # "context" | "resolve" | "by_name"
|
||||
resolve: Resolve | None
|
||||
wants_union: bool
|
||||
|
||||
def __init__(self, kind: str, resolve: Resolve | None = None, wants_union: bool = False) -> None:
|
||||
self.kind = kind
|
||||
self.resolve = resolve
|
||||
self.wants_union = wants_union
|
||||
|
||||
|
||||
class _ResolverPlan:
|
||||
"""A resolver's parameters and whether it is async, analyzed once."""
|
||||
|
||||
def __init__(self, fn: Callable[..., Any], params: dict[str, _ParamPlan], is_async: bool) -> None:
|
||||
self.fn = fn
|
||||
self.params = params
|
||||
self.is_async = is_async
|
||||
|
||||
|
||||
def _type_hints(fn: Callable[..., Any]) -> dict[str, Any]:
|
||||
"""Resolve type hints for a function or a callable object.
|
||||
|
||||
`typing.get_type_hints` raises on a callable *instance*; fall back to its
|
||||
`__call__`. Returns an empty mapping when hints cannot be resolved, matching
|
||||
`find_context_parameter`'s tolerance so callables without annotations (or with
|
||||
unresolvable ones) simply have no resolved parameters.
|
||||
"""
|
||||
target = fn if inspect.isroutine(fn) else getattr(type(fn), "__call__", fn)
|
||||
try:
|
||||
return typing.get_type_hints(target, include_extras=True)
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
|
||||
def _resolver_name(fn: Callable[..., Any]) -> str:
|
||||
"""Best-effort display name for error messages (callable objects lack `__name__`)."""
|
||||
return getattr(fn, "__name__", None) or type(fn).__name__
|
||||
|
||||
|
||||
def find_resolved_parameters(fn: Callable[..., Any]) -> dict[str, tuple[Resolve, bool]]:
|
||||
"""Find parameters of `fn` annotated `Annotated[_, Resolve(...)]`.
|
||||
|
||||
Returns a mapping of parameter name to `(Resolve, wants_union)`, where
|
||||
`wants_union` is True when the annotated type is an `ElicitationResult` member
|
||||
(the consumer wants the full outcome rather than the unwrapped model).
|
||||
"""
|
||||
hints = _type_hints(fn)
|
||||
resolved: dict[str, tuple[Resolve, bool]] = {}
|
||||
for name in inspect.signature(fn).parameters:
|
||||
annotation = hints.get(name)
|
||||
if get_origin(annotation) is not Annotated:
|
||||
# A `Resolve` marker is only honored at the top level; flag (rather than
|
||||
# silently drop) one buried in a union, e.g. `Annotated[T, Resolve(f)] | None`.
|
||||
if _contains_resolve(annotation):
|
||||
raise InvalidSignature(
|
||||
f"Parameter {name!r} of {_resolver_name(fn)!r} wraps `Resolve(...)` in a "
|
||||
"union; annotate the parameter directly as `Annotated[T, Resolve(...)]`"
|
||||
)
|
||||
continue
|
||||
type_arg, *metadata = get_args(annotation)
|
||||
marker = next((m for m in metadata if isinstance(m, Resolve)), None)
|
||||
if marker is not None:
|
||||
resolved[name] = (marker, _wants_union(type_arg))
|
||||
return resolved
|
||||
|
||||
|
||||
def _contains_resolve(annotation: Any) -> bool:
|
||||
"""True when a `Resolve` marker is nested inside `annotation` (e.g. a union member)."""
|
||||
if get_origin(annotation) is Annotated:
|
||||
return any(isinstance(m, Resolve) for m in get_args(annotation)[1:])
|
||||
return any(_contains_resolve(arg) for arg in get_args(annotation))
|
||||
|
||||
|
||||
def _wants_union(type_arg: Any) -> bool:
|
||||
"""True when `type_arg` is an `ElicitationResult` member (or a union of them).
|
||||
|
||||
Handles the subscripted `ElicitationResult[T]` alias (a `TypeAliasType` whose
|
||||
union is on the origin's `__value__`), the bare `ElicitationResult` alias (the
|
||||
`__value__` is on `type_arg` itself), an explicit `AcceptedElicitation[T] | ...`
|
||||
union, and a single member.
|
||||
"""
|
||||
# Unwrap the `ElicitationResult` alias whether it is bare or subscripted.
|
||||
value = getattr(type_arg, "__value__", None) or getattr(get_origin(type_arg), "__value__", None)
|
||||
if value is not None:
|
||||
type_arg = value
|
||||
members = get_args(type_arg) if get_origin(type_arg) is not None else (type_arg,)
|
||||
return any(isinstance(m, type) and issubclass(m, _ELICITATION_RESULT_MEMBERS) for m in members)
|
||||
|
||||
|
||||
def _resolver_key(fn: Callable[..., Any]) -> Hashable:
|
||||
"""Identity key for memoizing a resolver.
|
||||
|
||||
A bound method - pure-python (`inspect.ismethod`) or built-in (e.g. `obj.meth`
|
||||
on a C-extension type) - is recreated on each attribute access, so `id(fn)`
|
||||
differs every time. Key it by its underlying function (or name) plus its
|
||||
`__self__` identity so `auth.login` referenced in two places memoizes to one
|
||||
call. Everything else keys by `id`, so two distinct callables never collide
|
||||
even if they compare equal.
|
||||
"""
|
||||
bound_self = getattr(fn, "__self__", None)
|
||||
if bound_self is not None:
|
||||
# `__func__` (pure-python) has a stable identity; built-ins expose only a
|
||||
# stable `__name__`. Use the function's id or the name's value accordingly.
|
||||
func = getattr(fn, "__func__", None)
|
||||
underlying: Hashable = id(func) if func is not None else getattr(fn, "__name__", id(fn))
|
||||
return (underlying, id(bound_self))
|
||||
return id(fn)
|
||||
|
||||
|
||||
def build_resolver_plans(
|
||||
resolved_params: Mapping[str, tuple[Resolve, bool]],
|
||||
tool_arg_names: set[str],
|
||||
) -> dict[Hashable, _ResolverPlan]:
|
||||
"""Statically analyze the resolver DAG rooted at a tool's resolved parameters.
|
||||
|
||||
Raises:
|
||||
InvalidSignature: If a resolver has a cyclic dependency, or a resolver
|
||||
parameter cannot be classified (not a `Context`, a nested `Resolve`,
|
||||
or a tool argument by name).
|
||||
"""
|
||||
plans: dict[Hashable, _ResolverPlan] = {}
|
||||
|
||||
def analyze(fn: Callable[..., Any], stack: tuple[Hashable, ...]) -> None:
|
||||
key = _resolver_key(fn)
|
||||
if key in stack:
|
||||
raise InvalidSignature(f"Resolver {_resolver_name(fn)!r} has a cyclic dependency")
|
||||
if key in plans:
|
||||
return
|
||||
|
||||
hints = _type_hints(fn)
|
||||
sig = inspect.signature(fn)
|
||||
params: dict[str, _ParamPlan] = {}
|
||||
nested: list[Callable[..., Any]] = []
|
||||
for param_name in sig.parameters:
|
||||
annotation = hints.get(param_name)
|
||||
if annotation is not None and _is_context_annotation(annotation):
|
||||
params[param_name] = _ParamPlan("context")
|
||||
continue
|
||||
marker, wants_union = _resolve_marker(annotation)
|
||||
if marker is not None:
|
||||
params[param_name] = _ParamPlan("resolve", marker, wants_union)
|
||||
nested.append(marker.fn)
|
||||
continue
|
||||
if param_name in tool_arg_names:
|
||||
params[param_name] = _ParamPlan("by_name")
|
||||
continue
|
||||
raise InvalidSignature(
|
||||
f"Resolver {_resolver_name(fn)!r} parameter {param_name!r} cannot be resolved: "
|
||||
"expected a Context, an Annotated[_, Resolve(...)], or a tool argument by name"
|
||||
)
|
||||
|
||||
plans[key] = _ResolverPlan(fn, params, is_async_callable(fn))
|
||||
for dep in nested:
|
||||
analyze(dep, stack + (key,))
|
||||
|
||||
for marker, _ in resolved_params.values():
|
||||
analyze(marker.fn, ())
|
||||
return plans
|
||||
|
||||
|
||||
def _resolve_marker(annotation: Any) -> tuple[Resolve | None, bool]:
|
||||
if get_origin(annotation) is not Annotated:
|
||||
return None, False
|
||||
type_arg, *metadata = get_args(annotation)
|
||||
marker = next((m for m in metadata if isinstance(m, Resolve)), None)
|
||||
return marker, (_wants_union(type_arg) if marker is not None else False)
|
||||
|
||||
|
||||
def _is_context_annotation(annotation: Any) -> bool:
|
||||
if get_origin(annotation) is Annotated:
|
||||
annotation = get_args(annotation)[0]
|
||||
candidates = get_args(annotation) if get_origin(annotation) is not None else (annotation,)
|
||||
return any(isinstance(c, type) and issubclass(c, Context) for c in candidates)
|
||||
|
||||
|
||||
async def resolve_arguments(
|
||||
resolved_params: Mapping[str, tuple[Resolve, bool]],
|
||||
plans: Mapping[Hashable, _ResolverPlan],
|
||||
tool_args: Mapping[str, Any],
|
||||
context: Context[Any, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Resolve every `Resolve`-marked tool parameter into a concrete value.
|
||||
|
||||
Each resolver runs at most once (memoized by function identity). Returns a
|
||||
mapping of tool parameter name to the value to inject.
|
||||
|
||||
Raises:
|
||||
ToolError: If an elicited value is declined or cancelled and the consumer
|
||||
asked for the unwrapped model (rather than the result union).
|
||||
"""
|
||||
cache: dict[Hashable, ElicitationResult[Any]] = {}
|
||||
injected: dict[str, Any] = {}
|
||||
for name, (marker, wants_union) in resolved_params.items():
|
||||
outcome = await _resolve(marker.fn, plans, tool_args, context, cache)
|
||||
injected[name] = outcome if wants_union else _unwrap(outcome, name)
|
||||
return injected
|
||||
|
||||
|
||||
async def _resolve(
|
||||
fn: Callable[..., Any],
|
||||
plans: Mapping[Hashable, _ResolverPlan],
|
||||
tool_args: Mapping[str, Any],
|
||||
context: Context[Any, Any],
|
||||
cache: dict[Hashable, ElicitationResult[Any]],
|
||||
) -> ElicitationResult[Any]:
|
||||
key = _resolver_key(fn)
|
||||
if key in cache:
|
||||
return cache[key]
|
||||
|
||||
plan = plans[key]
|
||||
kwargs: dict[str, Any] = {}
|
||||
for param_name, param_plan in plan.params.items():
|
||||
if param_plan.kind == "context":
|
||||
kwargs[param_name] = context
|
||||
elif param_plan.kind == "by_name":
|
||||
kwargs[param_name] = tool_args[param_name]
|
||||
else:
|
||||
assert param_plan.resolve is not None
|
||||
dep_outcome = await _resolve(param_plan.resolve.fn, plans, tool_args, context, cache)
|
||||
kwargs[param_name] = dep_outcome if param_plan.wants_union else _unwrap(dep_outcome, param_name)
|
||||
|
||||
if plan.is_async:
|
||||
result = await fn(**kwargs)
|
||||
else:
|
||||
result = await anyio.to_thread.run_sync(lambda: fn(**kwargs))
|
||||
|
||||
outcome: ElicitationResult[Any]
|
||||
if isinstance(result, Elicit):
|
||||
elicit = cast("Elicit[BaseModel]", result)
|
||||
outcome = await context.elicit(elicit.message, elicit.schema)
|
||||
else:
|
||||
# A resolver may return any type (not just `BaseModel`); `model_construct`
|
||||
# wraps it as an accepted result without validating against the schema bound.
|
||||
outcome = cast("AcceptedElicitation[Any]", AcceptedElicitation.model_construct(data=result))
|
||||
|
||||
cache[key] = outcome
|
||||
return outcome
|
||||
|
||||
|
||||
def _unwrap(outcome: ElicitationResult[Any], name: str) -> Any:
|
||||
if isinstance(outcome, AcceptedElicitation):
|
||||
return outcome.data
|
||||
raise ToolError(f"Resolver for parameter {name!r} could not resolve: elicitation was {outcome.action}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Resolve",
|
||||
"Elicit",
|
||||
"ElicitationResult",
|
||||
"AcceptedElicitation",
|
||||
"DeclinedElicitation",
|
||||
"CancelledElicitation",
|
||||
"find_resolved_parameters",
|
||||
"build_resolver_plans",
|
||||
"resolve_arguments",
|
||||
]
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Hashable
|
||||
from functools import cached_property
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
|
|
@ -8,6 +8,11 @@ from mcp_types import Icon, ToolAnnotations
|
|||
from pydantic import BaseModel, Field
|
||||
|
||||
from mcp.server.mcpserver.exceptions import ToolError
|
||||
from mcp.server.mcpserver.resolve import (
|
||||
build_resolver_plans,
|
||||
find_resolved_parameters,
|
||||
resolve_arguments,
|
||||
)
|
||||
from mcp.server.mcpserver.utilities.context_injection import find_context_parameter
|
||||
from mcp.server.mcpserver.utilities.func_metadata import FuncMetadata, func_metadata
|
||||
from mcp.shared._callable_inspection import is_async_callable
|
||||
|
|
@ -32,6 +37,14 @@ class Tool(BaseModel):
|
|||
)
|
||||
is_async: bool = Field(description="Whether the tool is async")
|
||||
context_kwarg: str | None = Field(None, description="Name of the kwarg that should receive context")
|
||||
resolved_params: dict[str, Any] = Field(
|
||||
default_factory=lambda: {},
|
||||
exclude=True,
|
||||
description="Parameters filled by resolvers, mapped to (Resolve, wants_union)",
|
||||
)
|
||||
resolver_plans: dict[Hashable, Any] = Field(
|
||||
default_factory=lambda: {}, exclude=True, description="Static per-resolver parameter plans"
|
||||
)
|
||||
annotations: ToolAnnotations | None = Field(None, description="Optional annotations for the tool")
|
||||
icons: list[Icon] | None = Field(default=None, description="Optional list of icons for this tool")
|
||||
meta: dict[str, Any] | None = Field(default=None, description="Optional metadata for this tool")
|
||||
|
|
@ -67,13 +80,23 @@ class Tool(BaseModel):
|
|||
if context_kwarg is None: # pragma: no branch
|
||||
context_kwarg = find_context_parameter(fn)
|
||||
|
||||
resolved_params = find_resolved_parameters(fn)
|
||||
|
||||
skip_names = [context_kwarg] if context_kwarg is not None else []
|
||||
skip_names.extend(resolved_params)
|
||||
|
||||
func_arg_metadata = func_metadata(
|
||||
fn,
|
||||
skip_names=[context_kwarg] if context_kwarg is not None else [],
|
||||
skip_names=skip_names,
|
||||
structured_output=structured_output,
|
||||
)
|
||||
parameters = func_arg_metadata.arg_model.model_json_schema(by_alias=True)
|
||||
|
||||
# Match `model_dump_one_level`'s kwarg keys (alias when present, else field name)
|
||||
# so a by-name resolver param resolves to a key that exists at call time.
|
||||
tool_arg_names = {field.alias or name for name, field in func_arg_metadata.arg_model.model_fields.items()}
|
||||
resolver_plans = build_resolver_plans(resolved_params, tool_arg_names)
|
||||
|
||||
return cls(
|
||||
fn=fn,
|
||||
name=func_name,
|
||||
|
|
@ -83,6 +106,8 @@ class Tool(BaseModel):
|
|||
fn_metadata=func_arg_metadata,
|
||||
is_async=is_async,
|
||||
context_kwarg=context_kwarg,
|
||||
resolved_params=dict(resolved_params),
|
||||
resolver_plans=resolver_plans,
|
||||
annotations=annotations,
|
||||
icons=icons,
|
||||
meta=meta,
|
||||
|
|
@ -100,11 +125,26 @@ class Tool(BaseModel):
|
|||
ToolError: If the tool function raises during execution.
|
||||
"""
|
||||
try:
|
||||
pass_directly: dict[str, Any] = {}
|
||||
if self.context_kwarg is not None:
|
||||
pass_directly[self.context_kwarg] = context
|
||||
|
||||
# Resolvers see the same validated arguments the tool body receives:
|
||||
# validate once and reuse it, so a `default_factory`/stateful validator
|
||||
# can't hand a by-name resolver a different value than the body.
|
||||
pre_validated: dict[str, Any] | None = None
|
||||
if self.resolved_params:
|
||||
pre_validated = self.fn_metadata.validate_arguments(arguments)
|
||||
pass_directly |= await resolve_arguments(
|
||||
self.resolved_params, self.resolver_plans, pre_validated, context
|
||||
)
|
||||
|
||||
result = await self.fn_metadata.call_fn_with_arg_validation(
|
||||
self.fn,
|
||||
self.is_async,
|
||||
arguments,
|
||||
{self.context_kwarg: context} if self.context_kwarg is not None else None,
|
||||
pass_directly or None,
|
||||
pre_validated=pre_validated,
|
||||
)
|
||||
|
||||
if convert_result:
|
||||
|
|
|
|||
|
|
@ -69,21 +69,36 @@ class FuncMetadata(BaseModel):
|
|||
output_model: Annotated[type[BaseModel], WithJsonSchema(None)] | None = None
|
||||
wrap_output: bool = False
|
||||
|
||||
def validate_arguments(self, arguments_to_validate: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Validate raw arguments into a one-level kwargs dict (no function call).
|
||||
|
||||
Used to feed resolver dependency injection the validated tool arguments
|
||||
before the tool function itself runs.
|
||||
"""
|
||||
arguments_pre_parsed = self.pre_parse_json(arguments_to_validate)
|
||||
arguments_parsed_model = self.arg_model.model_validate(arguments_pre_parsed)
|
||||
return arguments_parsed_model.model_dump_one_level()
|
||||
|
||||
async def call_fn_with_arg_validation(
|
||||
self,
|
||||
fn: Callable[..., Any | Awaitable[Any]],
|
||||
fn_is_async: bool,
|
||||
arguments_to_validate: dict[str, Any],
|
||||
arguments_to_pass_directly: dict[str, Any] | None,
|
||||
pre_validated: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
"""Call the given function with arguments validated and injected.
|
||||
|
||||
Arguments are first attempted to be parsed from JSON, then validated against
|
||||
the argument model, before being passed to the function.
|
||||
the argument model, before being passed to the function. Pass `pre_validated`
|
||||
(the output of `validate_arguments`) to reuse an earlier validation pass -
|
||||
validating twice can re-run `default_factory`/stateful validators and hand the
|
||||
function different values than a caller already observed.
|
||||
"""
|
||||
arguments_pre_parsed = self.pre_parse_json(arguments_to_validate)
|
||||
arguments_parsed_model = self.arg_model.model_validate(arguments_pre_parsed)
|
||||
arguments_parsed_dict = arguments_parsed_model.model_dump_one_level()
|
||||
# Copy so a caller-provided `pre_validated` dict is never mutated in place.
|
||||
arguments_parsed_dict = dict(
|
||||
pre_validated if pre_validated is not None else self.validate_arguments(arguments_to_validate)
|
||||
)
|
||||
|
||||
arguments_parsed_dict |= arguments_to_pass_directly or {}
|
||||
|
||||
|
|
@ -150,7 +165,7 @@ class FuncMetadata(BaseModel):
|
|||
key_to_field_info[field_info.alias] = field_info
|
||||
|
||||
for data_key, data_value in data.items():
|
||||
if data_key not in key_to_field_info: # pragma: no cover
|
||||
if data_key not in key_to_field_info:
|
||||
continue
|
||||
|
||||
field_info = key_to_field_info[data_key]
|
||||
|
|
|
|||
Loading…
Reference in a new issue