Allow multiple models in temporal agent

[mattbrandman]

This pull request introduces enhanced support for model selection and management within Temporal agents, as well as improved handling of run context propagation. The main changes allow registering multiple models with a Temporal agent, selecting models by name or provider string at runtime (inside workflows), and ensuring the current run context is properly tracked across async boundaries. These improvements make it easier to use and configure multiple models in Temporal workflows, while maintaining safety and clarity in model selection.

Model selection and registration for Temporal agents:

• Added support for registering multiple models with a Temporal agent via the new additional_models argument, and for selecting a model by name or provider string at runtime within workflows. This includes validation to prevent duplicate or invalid model names and ensures that only registered models or provider strings can be selected during workflow execution. (pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py, [1] [2] [3] [4] [5] [6] [7] [8]

• Introduced the TemporalProviderFactory type and support for passing a provider factory to Temporal agents and models, enabling custom provider instantiation logic (e.g., injecting API keys from dependencies) when resolving models from provider strings. (pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py, [1] [2]; pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_model.py, [3] [4]

Model selection logic in Temporal model activities:

• Updated the Temporal model wrapper to support runtime model selection, including a context variable for the current model selection and logic to resolve the correct model instance for each request or stream activity. This ensures the correct model is used for each workflow step, whether by registered name or provider string. (pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_model.py, pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_model.pyR81-R112)

Run context propagation improvements:

• Introduced the CURRENT_RUN_CONTEXT context variable to track the current run context across asynchronous boundaries, and updated agent graph methods to set and reset this variable during model requests and streaming. This ensures that context-dependent logic (such as provider factories) has access to the correct run context throughout execution. (pydantic_ai_slim/pydantic_ai/_run_context.py, [1] [2]; pydantic_ai_slim/pydantic_ai/_agent_graph.py, [3] [4] [5] [6]

Other improvements and minor changes:

• Updated imports and type annotations to support the new features and improve clarity. (pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py, [1]; pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_model.py, [2] [3]

These changes collectively make Temporal agents more flexible, robust, and easier to configure for advanced use cases involving multiple models and dynamic provider selection.

[mattbrandman]

@DouweM let me know if this is more along the lines of what you are thinking. I did test this locally as well in our repo and it worked as expected.

[DouweM]

@mattbrandman Thanks for working on this Matt!

A few high level thoughts:

1. I wonder if we can use a single TemporalModel and just swap out what self.wrapped points at. We could pass the additional models to TemporalModel, and use a context manager + a model key pass into the request/request_stream activities to select the correct model
2. Activity names should never use generated values as users need to be able to keep already-running activities working when they change the underlying code, so if we need model-specific info in activity names, it should be the (normalized) provider:model name if a string is provided, and if an instance is provided the user should explicitly name it. So maybe have additional_models be dict[str, Model | str] | list[str], not allowing un-named model instances. Note that if point 1 works, we may not need the model name in the activity name at all, as they'll all use the same activities
3. If we do that, I suppose we can support arbitrary provider:model strs passed at agent.run time, as the TemporalModel can infer_model on them, and they don't need dedicated pre-registered activities. Specific model instances would still need to be registered up front, so they can be referenced by name. Maybe agent.run(model=...) should only take str | KnownModelName then, so we don't need to look up things by id().
4. I don't know if the extra generic param is worth the effort, as it restrict the types but not the specific instances or registered model names, so we'd need to rely on a runtime check and error anyway
5. If we support arbitrary model names on agent.run, it could be worth allowing a provider_factory to be registered to override how providers are built (see the arg on infer_model). Our version of that could also take run context, so the provider can be configured with an API key from deps, for example

[mattbrandman]

@DouweM changed the PR to be more inline with your comments above

[mattbrandman]

Confirmed this works locally. One thing that does appear to need updating but I'm not entirely sure where is that telemetry is printing the default model registered

[DouweM]

Model strings that will just be passed to infer_model don't really need to be pre-registered, do they? So we can probably drop this type from the union

[DouweM]

Is this also used on the model name the agent is initially created with? I would kind of expect it to

[DouweM]

I don't know if we need a whole ModelSelection object. I'd just store a single string, check if it's a key in the model instance map, and if not, call infer_model on it. But there may be another purpose to this that I'm missing

[DouweM]

The type already makes this impossible right? We typically don't check things the type checker would've caught already

[DouweM]

Instead of having a context var that's owned at this level and then passed into the TemporalModel, could TemporalModel own the context var and have a contextmanager method that we could use like with self._temporal_model.using_model('...'):?

[DouweM]

Probably not necessary; see above

[DouweM]

Since we don't edit this anymore, I'd rather move them back down where they were

[DouweM]

Unnecessary change?

[mattbrandman]

@DouweM made updates based on feedback.