Add takes_ctx arg to Tool.from_schema (#2615)

Author: dedeswim

docs/tools.md

@@ -471,7 +471,7 @@ _(This example is complete, it can be run "as is")_
 
 ### Custom Tool Schema
 
-If you have a function that lacks appropriate documentation (i.e. poorly named, no type information, poor docstring, use of \*args or \*\*kwargs and suchlike) then you can still turn it into a tool that can be effectively used by the agent with the [`Tool.from_schema`][pydantic_ai.Tool.from_schema] function. With this you provide the name, description and JSON schema for the function directly:
+If you have a function that lacks appropriate documentation (i.e. poorly named, no type information, poor docstring, use of \*args or \*\*kwargs and suchlike) then you can still turn it into a tool that can be effectively used by the agent with the [`Tool.from_schema`][pydantic_ai.Tool.from_schema] function. With this you provide the name, description, JSON schema, and whether the function takes a `RunContext` for the function directly:
 
 ```python
 from pydantic_ai import Agent, Tool
@@ -493,7 +493,8 @@ tool = Tool.from_schema(
         },
         'required': ['a', 'b'],
         'type': 'object',
-    }
+    },
+    takes_ctx=False,
 )
 
 test_model = TestModel()

pydantic_ai_slim/pydantic_ai/tools.py

@@ -266,6 +266,7 @@ def from_schema(
         name: str,
         description: str | None,
         json_schema: JsonSchemaValue,
+        takes_ctx: bool = False,
     ) -> Self:
         """Creates a Pydantic tool from a function and a JSON schema.
 
@@ -277,6 +278,8 @@ def from_schema(
             description: Used to tell the model how/when/why to use the tool.
                 You can provide few-shot examples as a part of the description.
             json_schema: The schema for the function arguments
+            takes_ctx: An optional boolean parameter indicating whether the function
+                accepts the context object as an argument.
 
         Returns:
             A Pydantic tool that calls the function
@@ -286,13 +289,13 @@ def from_schema(
             description=description,
             validator=SchemaValidator(schema=core_schema.any_schema()),
             json_schema=json_schema,
-            takes_ctx=False,
+            takes_ctx=takes_ctx,
             is_async=_utils.is_async_callable(function),
         )
 
         return cls(
             function,
-            takes_ctx=False,
+            takes_ctx=takes_ctx,
             name=name,
             description=description,
             function_schema=function_schema,

tests/models/cassettes/test_model_names/test_known_model_names.yaml

@@ -12,8 +12,6 @@ interactions:
     uri: https://us.inference.heroku.com/available-models
   response:
     headers:
-      connection:
-      - keep-alive
       content-length:
       - '760'
       content-security-policy:

tests/test_tools.py

@@ -1150,6 +1150,35 @@ def function(*args: Any, **kwargs: Any) -> str:
     assert agent._function_toolset.tools['foobar'].max_retries == 0
 
 
+def test_function_tool_from_schema_with_ctx():
+    def function(ctx: RunContext[str], *args: Any, **kwargs: Any) -> str:
+        assert len(args) == 0
+        assert set(kwargs) == {'one', 'two'}
+        return ctx.deps + 'I like being called like this'
+
+    json_schema = {
+        'type': 'object',
+        'additionalProperties': False,
+        'properties': {
+            'one': {'description': 'first argument', 'type': 'string'},
+            'two': {'description': 'second argument', 'type': 'object'},
+        },
+        'required': ['one', 'two'],
+    }
+    pydantic_tool = Tool[str].from_schema(
+        function, name='foobar', description='does foobar stuff', json_schema=json_schema, takes_ctx=True
+    )
+
+    assert pydantic_tool.takes_ctx is True
+    assert pydantic_tool.function_schema.takes_ctx is True
+
+    agent = Agent('test', tools=[pydantic_tool], retries=0, deps_type=str)
+    result = agent.run_sync('foobar', deps='Hello, ')
+    assert result.output == snapshot('{"foobar":"Hello, I like being called like this"}')
+    assert agent._function_toolset.tools['foobar'].takes_ctx is True
+    assert agent._function_toolset.tools['foobar'].max_retries == 0
+
+
 def test_function_tool_inconsistent_with_schema():
     def function(three: str, four: int) -> str:
         return 'Coverage made me call this'