Skip to content

[FEATURE] Allow opt-out from tool schema validation and normalization #2242

@charles-dyfis-net

Description

@charles-dyfis-net

Problem Statement

I have an existing codebase where Pydantic is used to generate tool-call schemas (as well as to parse tool responses). These schemas work well exactly as defined in the pre-strands codebase.

Strands breaks my tools: It adds unwanted generated Parameter ${name} values, it drops additionalProperties, and it otherwise fails to send content over the wire in its preexisting format with which my existing prompts have been carefully tested.

Proposed Solution

For callers using the existing strands tool decorator, it's reasonable to keep behavior exactly as it is -- we don't want a "fix" that breaks someone else's workflow.

A mix-in to opt out is a reasonable choice:

class PreValidatedTool:
    """A tool inheriting this class asserts that its tool_spec should be used exactly as defined"""

...and then later, in ToolRegistry, a gate around the relevant preexisting logic:

spec = copy.deepcopy(tool.tool_spec)
if not isinstance(tool, PreValidatedTool):
    spec = normalize_tool_spec(spec)
    self.validate_tool_spec(spec)
config[tool_name] = spec

Use Case

See "Problem Statement", for the most part. To provide a concrete example of how one might construct a new tool decorator that inherits from PreValidatedTool and passes through Pydantic schemas exactly as given:

from __future__ import annotations

import inspect
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, overload

from pydantic import TypeAdapter, ValidationError
from pydantic_core import PydanticSerializationError

from strands.interrupt import InterruptException
from strands.types._events import ToolInterruptEvent, ToolResultEvent
from strands.types.tools import AgentTool, ToolGenerator, ToolResultContent, ToolResultStatus, ToolSpec, ToolUse

VAR_KEYWORD = inspect.Parameter.VAR_KEYWORD

class PreValidatedTool:
    """Marker mixin: this tool's `tool_spec` is already well-formed."""

@dataclass
class RawToolResult:
    content: list[ToolResultContent]
    status: ToolResultStatus = "success"

class PydanticTool(AgentTool, PreValidatedTool):
    def __init__(self, cls: type, *, name: str | None = None, description: str | None = None) -> None:
        super().__init__()
        call = getattr(cls, "__call__", None)
        if not inspect.iscoroutinefunction(call):
            raise TypeError(f"{cls.__name__}.__call__ must be `async def`")

        sig = inspect.signature(call)
        if sig.return_annotation is inspect.Signature.empty:
            raise TypeError(f"{cls.__name__}.__call__ must declare a return annotation")

        self._input_adapter: TypeAdapter[Any] = TypeAdapter(cls)
        self._return_adapter: TypeAdapter[Any] = TypeAdapter(sig.return_annotation)

        self._tool_spec: ToolSpec = {
            "name": name or cls.__name__,
            "description": description or inspect.cleandoc(cls.__dict__.get("__doc__") or ""),
            "inputSchema": {"json": self._input_adapter.json_schema()},
        }

        self._has_var_keyword = any(p.kind is VAR_KEYWORD for p in sig.parameters.values())
        self._param_names = tuple(p.name for p in sig.parameters.values() if p.name != "self" and p.kind is not VAR_KEYWORD)

    @property
    def tool_name(self) -> str:
        return self._tool_spec["name"]

    @property
    def tool_spec(self) -> ToolSpec:
        return self._tool_spec

    @property
    def tool_type(self) -> str:
        return "function"

    async def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **_: Any) -> ToolGenerator:
        tool_use_id = tool_use.get("toolUseId", "unknown")

        def text_event(
            status: ToolResultStatus, text: str, exception: Exception | None = None
        ) -> ToolResultEvent:
            return ToolResultEvent({"toolUseId": tool_use_id, "status": status, "content": [{"text": text}]}, exception=exception)

        try:
            instance = self._input_adapter.validate_python(tool_use.get("input", {}) or {})
        except ValidationError as e:
            yield text_event("error", f"Error: Validation failed: {e}", e)
            return

        source = {**invocation_state, "invocation_state": invocation_state, "tool_use": tool_use}
        call_kwargs = (source if self._has_var_keyword else {n: source[n] for n in self._param_names if n in source})

        try:
            result = await instance(**call_kwargs)
        except InterruptException as e:
            yield ToolInterruptEvent(tool_use, [e.interrupt])
            return
        except Exception as e:
            yield text_event("error", f"Error: {type(e).__name__} - {e}", e)
            return

        if isinstance(result, RawToolResult):
            yield ToolResultEvent({"toolUseId": tool_use_id, "status": result.status, "content": result.content})
            return

        if isinstance(result, str):
            text = result
        else:
            try:
                text = self._return_adapter.dump_json(result).decode()
            except PydanticSerializationError:
                text = str(result)
        yield text_event("success", text)


@overload
def pydantic_tool(cls: type, /) -> PydanticTool: ...
@overload
def pydantic_tool(*, name: str | None = ..., description: str | None = ...) -> Callable[[type], PydanticTool]: ...
def pydantic_tool(cls: type | None = None, /, *, name: str | None = None, description: str | None = None) -> PydanticTool | Callable[[type], PydanticTool]:
    """Decorator/factory: `@pydantic_tool` or `@pydantic_tool(name=..., description=...)`."""
    if cls is not None:
        return PydanticTool(cls, name=name, description=description)
    return lambda c: PydanticTool(c, name=name, description=description)

(BTW, I'd be happy to submit a separate ticket/PR pair with exactly the above on its own).

Alternatives Solutions

No response

Additional Context

No response

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions