Repository: pydantic/pydantic-ai
File: tests/typed_agent.py (Python)

Code before:
Tool(foobar_ctx, takes_ctx=True)
Tool(foobar_ctx)
Tool(foobar_plain, takes_ctx=False)
assert_type(Tool(foobar_plain), Tool[object])
assert_type(Tool(foobar_plain), Tool)

# unfortunately we can't type check these cases, since from a typing perspect `foobar_ctx` is valid as a plain tool
Tool(foobar_ctx, takes_ctx=False)
Tool(foobar_plain, takes_ctx=True)

Agent('test', tools=[foobar_ctx], deps_type=int)
Agent('test', tools=[foobar_plain], deps_type=int)
Agent('test', tools=[foobar_plain])
Agent('test', tools=[Tool(foobar_ctx)], deps_type=int)
Agent('test', tools=[Tool(foobar_ctx), foobar_ctx, foobar_plain], deps_type=int)
Agent('test', tools=[Tool(foobar_ctx), foobar_ctx, Tool(foobar_plain)], deps_type=int)

Agent('test', tools=[foobar_ctx], deps_type=str)  # pyright: ignore[reportArgumentType,reportCallIssue]
Agent('test', tools=[Tool(foobar_ctx), Tool(foobar_plain)], deps_type=str)  # pyright: ignore[reportArgumentType,reportCallIssue]
Agent('test', tools=[foobar_ctx])  # pyright: ignore[reportArgumentType,reportCallIssue]
Agent('test', tools=[Tool(foobar_ctx)])  # pyright: ignore[reportArgumentType,reportCallIssue]
# since deps are not set, they default to `None`, so can't be `int`
Agent('test', tools=[Tool(foobar_plain)], deps_type=int)  # pyright: ignore[reportArgumentType,reportCallIssue]

# TextOutput with RunContext
Agent('test', output_type=TextOutput(str_to_regex_with_ctx), deps_type=int)
Agent('test', output_type=TextOutput(str_to_regex_with_ctx), deps_type=str)  # pyright: ignore[reportArgumentType,reportCallIssue]
Agent('test', output_type=TextOutput(str_to_regex_with_ctx))  # pyright: ignore[reportArgumentType,reportCallIssue]

# prepare example from docs:


def greet(name: str) -> str:
    return f'hello {name}'


async def prepare_greet(ctx: RunContext[str], tool_def: ToolDefinition) -> ToolDefinition | None:
    d = f'Name of the {ctx.deps} to greet.'
    tool_def.parameters_json_schema['properties']['name']['description'] = d
    return tool_def


greet_tool = Tool(greet, prepare=prepare_greet)
assert_type(greet_tool, Tool[str])
greet_agent = Agent[str, str]('test', tools=[greet_tool], deps_type=str)

result = greet_agent.run_sync('testing...', deps='human')
assert result.output == '{"greet":"hello a"}'

if not MYPY:
    default_agent = Agent()

Review comment (suggestion):
To verify that `contravariant=True` is the right option, we should also have an example where the `deps_type` is a subclass of the type on the function's ctx, e.g. `bool`. That should work, because a function that takes an int is also able to handle a bool (which is essentially 1 or 0.

Code after:


Tool(foobar_ctx, takes_ctx=True)
Tool(foobar_ctx)
Tool(foobar_plain, takes_ctx=False)
assert_type(Tool(foobar_plain), Tool[object])
assert_type(Tool(foobar_plain), Tool)

# unfortunately we can't type check these cases, since from a typing perspect `foobar_ctx` is valid as a plain tool
Tool(foobar_ctx, takes_ctx=False)
Tool(foobar_plain, takes_ctx=True)

Agent('test', tools=[foobar_ctx], deps_type=int)
Agent('test', tools=[foobar_plain], deps_type=int)
Agent('test', tools=[foobar_plain])
Agent('test', tools=[Tool(foobar_ctx)], deps_type=int)
Agent('test', tools=[Tool(foobar_ctx), foobar_ctx, foobar_plain], deps_type=int)
Agent('test', tools=[Tool(foobar_ctx), foobar_ctx, Tool(foobar_plain)], deps_type=int)

Agent('test', tools=[foobar_ctx], deps_type=str)  # pyright: ignore[reportArgumentType,reportCallIssue]
Agent('test', tools=[Tool(foobar_ctx), Tool(foobar_plain)], deps_type=str)  # pyright: ignore[reportArgumentType,reportCallIssue]
Agent('test', tools=[foobar_ctx])  # pyright: ignore[reportArgumentType,reportCallIssue]
Agent('test', tools=[Tool(foobar_ctx)])  # pyright: ignore[reportArgumentType,reportCallIssue]
# since deps are not set, they default to `None`, so can't be `int`
Agent('test', tools=[Tool(foobar_plain)], deps_type=int)  # pyright: ignore[reportArgumentType,reportCallIssue]

# TextOutput with RunContext uses RunContext[Any], so deps_type is not checked.
# This is intentional: type checking deps in output functions isn't feasible because
# ToolOutput and plain output functions take arbitrary args, so the type checker
# treats RunContext as just another arg rather than enforcing deps_type compatibility.
text_output_with_ctx = TextOutput(str_to_regex_with_ctx)
assert_type(text_output_with_ctx, TextOutput[re.Pattern[str]])
Agent('test', output_type=text_output_with_ctx, deps_type=int)
Agent('test', output_type=text_output_with_ctx, deps_type=str)
Agent('test', output_type=text_output_with_ctx)

# prepare example from docs:


def greet(name: str) -> str:
    return f'hello {name}'


async def prepare_greet(ctx: RunContext[str], tool_def: ToolDefinition) -> ToolDefinition | None:
    d = f'Name of the {ctx.deps} to greet.'
    tool_def.parameters_json_schema['properties']['name']['description'] = d
    return tool_def


greet_tool = Tool(greet, prepare=prepare_greet)
assert_type(greet_tool, Tool[str])