[Draft] Add input model to Agent

[strawgate]

re: #2329

Having it be properly typed from the agent init -> run -> runcontext means this is not a small change :(

[DouweM]

@strawgate Thanks for looking into this!

I still like the idea conceptually, but the number of places the new generic parameter now needs to be accounted for gives me pause, so I don't think we should do this unless it clearly enables new usage patterns and workarounds are clearly insufficient. I feel like how complex a class feels scales much faster than linearly with the number of generic parameters, so 3 is much worse than 2, not just 50% worse, especially since most users likely won't actually need to use this, but would see it in their IDE etc.

Looking back at your issue:

Placing required inputs in deps_type is awkward as ideally you'd share dependencies but inputs are not often shared when running multi agent

I feel like it would work almost as well to make the deps type itself a dataclass generic in the input type, and then use dataclasses.replace at agent-run time to inject the input object into the shared deps object, likely with some typing.cast to go from MyDeps[Any] to MyDeps[MyInput].

I was also thinking of this feature in context of handoffs where the parent agent will need to know what type its sub-agent expects, but that could reasonably be worked around with something like handoffs=[Handoff(agent, input_type=MyInput, input_serializer=format_as_xml)].

I could also imagine a new Agent subclass that takes an input object in its run methods instead of user_prompt, with a convenience method like handoffs=[agent.with_input(MyInput, format_as_xml)] to wrap the current agent in the input-taking one. (We're already adding an AbstractAgent and WrapperAgent in #2225).

The other concern is that if we add this to the main Agent, it's not clear what should happen if there is no dynamic instructions function that actually uses the inpu1. Should we automatically serialize it and put it in the user prompt? Or have a default dynamic instructions function? That'd be less of an issue if using an input model required some explicit wrapping, either of the run method (with a function that injects it into deps as I suggested above) or the agent as a whole.

Curious to hear your thoughts. It'd be great to see specific examples of use cases that are hard-to-impossible to implement using one of these workarounds, to convince ourselves it's worth the extra generic param.

[strawgate]

@strawgate Thanks for looking into this!

I still like the idea conceptually, but the number of places the new generic parameter now needs to be accounted for gives me pause, so I don't think we should do this unless it clearly enables new usage patterns and workarounds are clearly insufficient. I feel like how complex a class feels scales much faster than linearly with the number of generic parameters, so 3 is much worse than 2, not just 50% worse, especially since most users likely won't actually need to use this, but would see it in their IDE etc.

I felt the same way making the PR, not only that but the ordering of the generics is a bit unfortunate too....

I was also thinking of this feature in context of handoffs where the parent agent will need to know what type its sub-agent expects, but that could reasonably be worked around with something like handoffs=[Handoff(agent, input_type=MyInput, input_serializer=format_as_xml)].

Handoffs are a good example of a use-case, which is basically just another way of exposing the agent as a tool. Unfortunately, I think having the caller of the agent determine the input_type puts the information in the wrong place.

I think the key items here are 1) availability in the runcontext (not just compilation into instructions) and 2) the ability to provide the model to the agent in lieu of a user_prompt.

I think given that the only viable path forward is probably to see if we can embed it in the depstype, but that's likely also a big change as you mention -- there's no bounds on deps_type today (it can be a str or Path, etc).

I'll noodle on this a bit and maybe experiment with your deps_type suggestion.

[DouweM]

were you suggesting changing deps_type to a data class generic? or something more specific to input_model/input_type?

@strawgate Responding to your question from Slack here to keep the context in one place: I was actually not thinking of changing anything in Pydantic AI to do that, just whether in your code a workaround could be to have your own (shared) deps type be generic in the input type of a particular agent. But I don't expect that to actually be sufficient.

Here's another idea: In #2225, we'll be introducing an AbstractAgent superclass to Agent, and a WrapperAgent subclass similar to WrapperModel/Toolset. I could imagine an InputAgent(agent, input_type: type[InputT]) subclass of WrapperAgent that requires the wrapped agent to have a deps_type that's a dataclass with an input field (similar to https://ai.pydantic.dev/ag-ui/#state-management). That could then override run (etc) to have an additional input: InputT argument alongside deps: DepsT and inject it using replace(deps, input=input). As the run methods take many arguments they shouldn't be directly used as tool functions (as the model will try to fill in stuff it shouldn't), so this InputAgent could also have a new run_with_input(input: InputT) (or as_tool or something) method that calls run with just that argument. To be able to bind input-less deps as well, either something ugly like run_with_deps_and_input(deps, input) that can be used as partial(input_agent.run_with_deps_and_input, deps) to get a function that just takes input, or something like input_agent.with_deps(deps).run_with_input if with_deps returns a new InputAgentWithDeps that stores the deps and uses them in all run methods.

I don't quite like that yet and I'm sure it'd hit some hurdles, but I think it's worth experimenting with whether a new Agent class could help us out here without cluttering the main one.