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

Code before:

# 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
text_output_with_ctx = TextOutput(str_to_regex_with_ctx)
assert_type(text_output_with_ctx, TextOutput[re.Pattern[str], int])
Agent('test', output_type=text_output_with_ctx, deps_type=int)
Agent('test', output_type=text_output_with_ctx, deps_type=bool)  # bool is subclass of int, works with contravariant
# NOTE: The following don't produce type errors because _OutputSpecItem uses TextOutput[T_co, Any]
# which erases the deps type constraint.
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])
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):
If we can't make those give a typing error, then maybe it's not worth adding `TextOutputAgentDepsT` and we can just make it `RunContext[Any]`. 

Ideally, we would go a step further and add `DepsT` to `_OutputSpecItem`, `OutputSpec` etc, so that the user _is_ warned if their `ctx: RunContext[Foo]` arg doesn't match the agent's `deps_type`, which would work for `TextOutput(func)`, but unfortunately there's now way to make that work for output functions in general, e.g. `output_type=func` or `output_type=ToolOutput(func)`, as those functions that arbitrary args, so if the first arg `ctx: RunContext[Foo]` while `deps_type=Bar`, the type checker will just treat it like an arbitrary arg, instead of complaining. (ping me if that doesn't make sense; typing is tricky)

So I think the best we can do is just make it `RunContext[Any]` and accept that this part isn't type checked.

Code after:

# 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])
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()