Python: Added explicit schema handling to @tool decorator (#3734)

* Added explicit schema handling to @tool decorator

* Resolved comments

python/packages/core/agent_framework/_tools.py

@@ -1239,6 +1239,7 @@ def tool(
     *,
     name: str | None = None,
     description: str | None = None,
+    schema: type[BaseModel] | Mapping[str, Any] | None = None,
     approval_mode: Literal["always_require", "never_require"] | None = None,
     max_invocations: int | None = None,
     max_invocation_exceptions: int | None = None,
@@ -1252,6 +1253,7 @@ def tool(
     *,
     name: str | None = None,
     description: str | None = None,
+    schema: type[BaseModel] | Mapping[str, Any] | None = None,
     approval_mode: Literal["always_require", "never_require"] | None = None,
     max_invocations: int | None = None,
     max_invocation_exceptions: int | None = None,
@@ -1264,6 +1266,7 @@ def tool(
     *,
     name: str | None = None,
     description: str | None = None,
+    schema: type[BaseModel] | Mapping[str, Any] | None = None,
     approval_mode: Literal["always_require", "never_require"] | None = None,
     max_invocations: int | None = None,
     max_invocation_exceptions: int | None = None,
@@ -1279,6 +1282,9 @@ def tool(
     with a string description as the second argument. You can also use Pydantic's
     ``Field`` class for more advanced configuration.
 
+    Alternatively, you can provide an explicit schema via the ``schema`` parameter
+    to bypass automatic inference from the function signature.
+
     Args:
         func: The function to decorate.
 
@@ -1287,6 +1293,13 @@ def tool(
             attribute will be used.
         description: A description of the function. If not provided, the function's
             docstring will be used.
+        schema: An explicit input schema for the function. This can be a Pydantic
+            ``BaseModel`` subclass or a JSON schema dictionary (``Mapping[str, Any]``).
+            When a dictionary is provided, it must be a flat object schema with a
+            ``properties`` key (complex JSON Schema features such as ``oneOf``,
+            ``$ref``, or nested compositions are not supported).
+            When provided, the schema is used instead of inferring one from the
+            function's signature. Defaults to ``None`` (infer from signature).
         approval_mode: Whether or not approval is required to run this tool.
             Default is that approval is required.
         max_invocations: The maximum number of times this function can be invoked.
@@ -1341,6 +1354,21 @@ def tool(
                 # Simulate async operation
                 return f"Weather in {location}"
 
+
+            # With an explicit Pydantic model schema
+            from pydantic import BaseModel, Field
+
+
+            class WeatherInput(BaseModel):
+                location: Annotated[str, Field(description="City name")]
+                unit: str = "celsius"
+
+
+            @tool(schema=WeatherInput)
+            def get_weather(location: str, unit: str = "celsius") -> str:
+                '''Get weather for a location.'''
+                return f"Weather in {location}: 22 {unit}"
+
     """
 
     def decorator(func: Callable[..., ReturnT | Awaitable[ReturnT]]) -> FunctionTool[Any, ReturnT]:
@@ -1356,6 +1384,7 @@ def tool(
                 max_invocation_exceptions=max_invocation_exceptions,
                 additional_properties=additional_properties or {},
                 func=f,
+                input_model=schema,
             )
 
         return wrapper(func)

python/packages/core/tests/core/test_tools.py

@@ -70,6 +70,102 @@ def test_tool_decorator_without_args():
     assert test_tool.approval_mode == "never_require"
 
 
+def test_tool_decorator_with_pydantic_schema():
+    """Test that the tool decorator accepts an explicit Pydantic model schema."""
+    from pydantic import Field
+
+    class MyInput(BaseModel):
+        location: Annotated[str, Field(description="City name")]
+        unit: str = "celsius"
+
+    @tool(name="weather", description="Get weather", schema=MyInput)
+    def get_weather(location: str, unit: str = "celsius") -> str:
+        return f"{location}: {unit}"
+
+    assert isinstance(get_weather, FunctionTool)
+    assert get_weather.name == "weather"
+    params = get_weather.parameters()
+    assert "location" in params["properties"]
+    assert params["properties"]["location"].get("description") == "City name"
+    assert get_weather("Seattle") == "Seattle: celsius"
+    assert get_weather("Seattle", "fahrenheit") == "Seattle: fahrenheit"
+
+
+def test_tool_decorator_with_json_schema_dict():
+    """Test that the tool decorator accepts an explicit JSON schema dict."""
+
+    json_schema = {
+        "type": "object",
+        "properties": {
+            "query": {"type": "string", "description": "Search query"},
+            "max_results": {"type": "integer", "default": 10},
+        },
+        "required": ["query"],
+    }
+
+    @tool(name="search", description="Search tool", schema=json_schema)
+    def search(query: str, max_results: int = 10) -> str:
+        return f"Searching for: {query} (max {max_results})"
+
+    assert isinstance(search, FunctionTool)
+    params = search.parameters()
+    assert params["properties"]["query"]["type"] == "string"
+    assert params["properties"]["query"]["description"] == "Search query"
+    assert "max_results" in params["properties"]
+    assert search("hello") == "Searching for: hello (max 10)"
+
+
+def test_tool_decorator_schema_none_default():
+    """Test that schema=None (default) still infers from function signature."""
+
+    @tool(name="adder", schema=None)
+    def add(x: int, y: int) -> int:
+        return x + y
+
+    assert isinstance(add, FunctionTool)
+    params = add.parameters()
+    assert params == {
+        "properties": {"x": {"title": "X", "type": "integer"}, "y": {"title": "Y", "type": "integer"}},
+        "required": ["x", "y"],
+        "title": "adder_input",
+        "type": "object",
+    }
+    assert add(1, 2) == 3
+
+
+async def test_tool_decorator_with_schema_invoke():
+    """Test that invoke works correctly with explicit schema."""
+
+    class CalcInput(BaseModel):
+        a: int
+        b: int
+
+    @tool(name="calc", description="Calculator", schema=CalcInput)
+    def calculate(a: int, b: int) -> int:
+        return a + b
+
+    result = await calculate.invoke(arguments=CalcInput(a=3, b=7))
+    assert result == 10
+
+
+def test_tool_decorator_with_schema_overrides_annotations():
+    """Test that explicit schema completely overrides function signature inference."""
+    from pydantic import Field
+
+    class DetailedInput(BaseModel):
+        location: Annotated[str, Field(description="The city and state")]
+        unit: Annotated[str, Field(description="Temperature unit")] = "celsius"
+
+    @tool(schema=DetailedInput)
+    def get_weather(location: str, unit: str = "celsius") -> str:
+        """Get weather for a location."""
+        return f"{location}: {unit}"
+
+    params = get_weather.parameters()
+    assert params["properties"]["location"].get("description") == "The city and state"
+    assert params["properties"]["unit"].get("description") == "Temperature unit"
+
+
 def test_tool_without_args():
     """Test the tool decorator."""
 

python/samples/getting_started/tools/README.md

@@ -19,6 +19,7 @@ keep `approval_mode="always_require"` unless you are confident in the tool behav
 | [`function_tool_with_thread_injection.py`](function_tool_with_thread_injection.py) | Shows how to access the current `thread` object inside a local tool via `**kwargs`. |
 | [`function_tool_with_max_exceptions.py`](function_tool_with_max_exceptions.py) | Shows how to limit the number of times a tool can fail with exceptions using `max_invocation_exceptions`. Useful for preventing expensive tools from being called repeatedly when they keep failing. |
 | [`function_tool_with_max_invocations.py`](function_tool_with_max_invocations.py) | Demonstrates limiting the total number of times a tool can be invoked using `max_invocations`. Useful for rate-limiting expensive operations or ensuring tools are only called a specific number of times per conversation. |
+| [`function_tool_with_explicit_schema.py`](function_tool_with_explicit_schema.py) | Demonstrates how to provide an explicit Pydantic model or JSON schema dictionary to the `@tool` decorator via the `schema` parameter, bypassing automatic inference from the function signature. |
 | [`tool_in_class.py`](tool_in_class.py) | Shows how to use the `tool` decorator with class methods to create stateful tools. Demonstrates how class state can control tool behavior dynamically, allowing you to adjust tool functionality at runtime by modifying class properties. |
 
 ## Key Concepts
@@ -26,6 +27,7 @@ keep `approval_mode="always_require"` unless you are confident in the tool behav
 ### Local Tool Features
 
 - **Function Declarations**: Define tool schemas without implementations for testing or external tools
+- **Explicit Schema**: Provide a Pydantic model or JSON schema dict to control the tool's parameter schema directly
 - **Dependency Injection**: Create tools from configurations with runtime-injected implementations
 - **Error Handling**: Gracefully handle and recover from tool execution failures
 - **Approval Workflows**: Require user approval before executing sensitive or important operations
@@ -55,6 +57,23 @@ def sensitive_operation(data: Annotated[str, "Data to process"]) -> str:
     return f"Processed: {data}"
 ```
 
+#### Tool with Explicit Schema
+
+```python
+from pydantic import BaseModel, Field
+from agent_framework import tool
+from typing import Annotated
+
+class WeatherInput(BaseModel):
+    location: Annotated[str, Field(description="City name")]
+    unit: str = "celsius"
+
+@tool(schema=WeatherInput)
+def get_weather(location: str, unit: str = "celsius") -> str:
+    """Get the weather for a location."""
+    return f"Weather in {location}: 22 {unit}"
+```
+
 #### Tool with Invocation Limits
 
 ```python

python/samples/getting_started/tools/function_tool_with_explicit_schema.py

@@ -0,0 +1,81 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+Function Tool with Explicit Schema Example
+
+This example demonstrates how to provide an explicit schema to the @tool decorator
+using the `schema` parameter, bypassing the automatic inference from the function
+signature. This is useful when you want full control over the tool's parameter
+schema that the AI model sees, or when the function signature does not accurately
+represent the desired schema.
+
+Two approaches are shown:
+1. Using a Pydantic BaseModel subclass as the schema
+2. Using a raw JSON schema dictionary as the schema
+"""
+
+import asyncio
+from typing import Annotated
+
+from agent_framework import tool
+from agent_framework.openai import OpenAIResponsesClient
+from pydantic import BaseModel, Field
+
+
+# Approach 1: Pydantic model as explicit schema
+class WeatherInput(BaseModel):
+    """Input schema for the weather tool."""
+
+    location: Annotated[str, Field(description="The city name to get weather for")]
+    unit: Annotated[str, Field(description="Temperature unit: celsius or fahrenheit")] = "celsius"
+
+
+@tool(
+    name="get_weather",
+    description="Get the current weather for a given location.",
+    schema=WeatherInput,
+    approval_mode="never_require",
+)
+def get_weather(location: str, unit: str = "celsius") -> str:
+    """Get the current weather for a location."""
+    return f"The weather in {location} is 22 degrees {unit}."
+
+
+# Approach 2: JSON schema dictionary as explicit schema
+get_current_time_schema = {
+    "type": "object",
+    "properties": {
+        "timezone": {"type": "string", "description": "The timezone to get the current time for", "default": "UTC"},
+    },
+}
+
+
+@tool(
+    name="get_current_time",
+    description="Get the current time in a given timezone.",
+    schema=get_current_time_schema,
+    approval_mode="never_require",
+)
+def get_current_time(timezone: str = "UTC") -> str:
+    """Get the current time."""
+    from datetime import datetime
+    from zoneinfo import ZoneInfo
+
+    return f"The current time in {timezone} is {datetime.now(ZoneInfo(timezone)).isoformat()}"
+
+
+async def main():
+    agent = OpenAIResponsesClient().as_agent(
+        name="AssistantAgent",
+        instructions="You are a helpful assistant. Use the available tools to answer questions.",
+        tools=[get_weather, get_current_time],
+    )
+
+    query = "What is the weather in Seattle and what time is it?"
+    print(f"User: {query}")
+    result = await agent.run(query)
+    print(f"Result: {result.text}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())