fixes types

Author: mattbrandman

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py

@@ -39,7 +39,7 @@
     ToolFuncEither,
 )
 
-from ._model import ModelSelection, TemporalModel, TemporalProviderFactory
+from ._model import TemporalModel, TemporalProviderFactory
 from ._run_context import TemporalRunContext
 from ._toolset import TemporalWrapperToolset, temporalize_toolset
 
@@ -124,11 +124,6 @@ def __init__(
             )
 
         self._registered_model_instances: dict[str, Model] = {'default': wrapped_model}
-        self._default_selection = ModelSelection(model_key='default', model_name=None)
-        self._model_selection_var: ContextVar[ModelSelection] = ContextVar(
-            '_temporal_model_selection',
-            default=self._default_selection,
-        )
 
         # start_to_close_timeout is required
         activity_config = activity_config or ActivityConfig(start_to_close_timeout=timedelta(seconds=60))
@@ -182,7 +177,6 @@ async def streamed_response():
             deps_type=self.deps_type,
             run_context_type=self.run_context_type,
             event_stream_handler=self.event_stream_handler,
-            model_selection_var=self._model_selection_var,
             model_instances=self._registered_model_instances,
             provider_factory=provider_factory,
         )
@@ -229,24 +223,22 @@ def _normalize_model_name(self, name: str) -> str:
             raise UserError(f'Duplicate model name {normalized!r} provided to `additional_models`.')
         return normalized
 
-    def _select_model(self, model: models.Model | models.KnownModelName | str | None = None) -> ModelSelection:
-        """Select the appropriate model based on the runtime parameter."""
-        if model is None:
-            return self._default_selection
+    def _select_model(self, model: models.Model | models.KnownModelName | str | None = None) -> str | None:
+        """Select the appropriate model based on the runtime parameter.
+
+        Returns a string that will be checked against registered model instances,
+        or passed to infer_model if not found. Returns None to use the default model.
+        """
+        if model is None or model == 'default':
+            return None
 
         if isinstance(model, Model):
             raise UserError(
                 'Model instances cannot be selected at runtime inside a Temporal workflow. '
                 'Register the model via the mapping form of `additional_models` and reference it by name.'
             )
 
-        if model == 'default':
-            return self._default_selection
-
-        if model in self._registered_model_instances:
-            return ModelSelection(model_key=model, model_name=None)
-
-        return ModelSelection(model_key=None, model_name=model)
+        return model
 
     @property
     def name(self) -> str | None:
@@ -299,11 +291,13 @@ def temporal_activities(self) -> list[Callable[..., Any]]:
         return self._temporal_activities
 
     @contextmanager
-    def _temporal_overrides(self, selection: ModelSelection | None = None) -> Iterator[None]:
+    def _temporal_overrides(self, selection: str | None = None) -> Iterator[None]:
         # We reset tools here as the temporalized function toolset is already in self._toolsets.
-        with super().override(model=self._temporal_model, toolsets=self._toolsets, tools=[]):
+        with (
+            super().override(model=self._temporal_model, toolsets=self._toolsets, tools=[]),
+            self._temporal_model.using_model(selection),
+        ):
             token = self._temporal_overrides_active.set(True)
-            selection_token = self._model_selection_var.set(selection or self._default_selection)
             try:
                 yield
             except PydanticSerializationError as e:
@@ -312,7 +306,6 @@ def _temporal_overrides(self, selection: ModelSelection | None = None) -> Iterat
                 ) from e
             finally:
                 self._temporal_overrides_active.reset(token)
-                self._model_selection_var.reset(selection_token)
 
     @overload
     async def run(
@@ -416,12 +409,12 @@ async def main():
                 'Event stream handler cannot be set at agent run time inside a Temporal workflow, it must be set at agent creation time.'
             )
 
-        selection = self._default_selection
+        selection: str | None = None
         if workflow.in_workflow() and model is not None:
             selection = self._select_model(model)
             model = None
 
-        with self._temporal_overrides(selection if workflow.in_workflow() else self._default_selection):
+        with self._temporal_overrides(selection if workflow.in_workflow() else None):
             return await super().run(
                 user_prompt,
                 output_type=output_type,

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_model.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from collections.abc import AsyncIterator, Callable, Mapping
-from contextlib import asynccontextmanager
+from collections.abc import AsyncIterator, Callable, Iterator, Mapping
+from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
 from datetime import datetime
@@ -33,14 +33,7 @@ class _RequestParams:
     model_settings: dict[str, Any] | None
     model_request_parameters: ModelRequestParameters
     serialized_run_context: Any | None = None
-    model_key: str | None = None
-    model_name: str | None = None
-
-
-@dataclass(frozen=True)
-class ModelSelection:
-    model_key: str | None = None
-    model_name: str | None = None
+    model_selection: str | None = None
 
 
 TemporalProviderFactory = Callable[[str, RunContext[Any] | None, Any | None], Provider[Any]]
@@ -85,22 +78,18 @@ def __init__(
         deps_type: type[AgentDepsT],
         run_context_type: type[TemporalRunContext[AgentDepsT]] = TemporalRunContext[AgentDepsT],
         event_stream_handler: EventStreamHandler[Any] | None = None,
-        model_selection_var: ContextVar[ModelSelection] | None = None,
         model_instances: Mapping[str, Model] | None = None,
         provider_factory: TemporalProviderFactory | None = None,
     ):
         super().__init__(model)
         self.activity_config = activity_config
         self.run_context_type = run_context_type
         self.event_stream_handler = event_stream_handler
-        self._model_selection_var = model_selection_var
+        self._model_selection_var: ContextVar[str | None] = ContextVar('_temporal_model_selection', default=None)
         self._model_instances = dict(model_instances or {})
         self._provider_factory = provider_factory
 
-        request_activity_name = f'{activity_name_prefix}__model_request'
-        request_stream_activity_name = f'{activity_name_prefix}__model_request_stream'
-
-        @activity.defn(name=request_activity_name)
+        @activity.defn(name=f'{activity_name_prefix}__model_request')
         async def request_activity(params: _RequestParams, deps: Any) -> ModelResponse:
             model_for_request = self._resolve_model(params, deps)
             return await model_for_request.request(
@@ -112,7 +101,7 @@ async def request_activity(params: _RequestParams, deps: Any) -> ModelResponse:
         self.request_activity = request_activity
         self.request_activity.__annotations__['deps'] = deps_type
 
-        async def request_stream_activity(params: _RequestParams, deps: Any) -> ModelResponse:
+        async def request_stream_activity(params: _RequestParams, deps: AgentDepsT) -> ModelResponse:
             # An error is raised in `request_stream` if no `event_stream_handler` is set.
             assert self.event_stream_handler is not None
 
@@ -135,7 +124,9 @@ async def request_stream_activity(params: _RequestParams, deps: Any) -> ModelRes
         # Set type hint explicitly so that Temporal can take care of serialization and deserialization
         request_stream_activity.__annotations__['deps'] = deps_type
 
-        self.request_stream_activity = activity.defn(name=request_stream_activity_name)(request_stream_activity)
+        self.request_stream_activity = activity.defn(name=f'{activity_name_prefix}__model_request_stream')(
+            request_stream_activity
+        )
 
     @property
     def temporal_activities(self) -> list[Callable[..., Any]]:
@@ -168,8 +159,7 @@ async def request(
                     model_settings=cast(dict[str, Any] | None, model_settings),
                     model_request_parameters=model_request_parameters,
                     serialized_run_context=serialized_run_context,
-                    model_key=selection.model_key,
-                    model_name=selection.model_name,
+                    model_selection=selection,
                 ),
                 deps,
             ],
@@ -212,8 +202,7 @@ async def request_stream(
                     model_settings=cast(dict[str, Any] | None, model_settings),
                     model_request_parameters=model_request_parameters,
                     serialized_run_context=serialized_run_context,
-                    model_key=selection.model_key,
-                    model_name=selection.model_name,
+                    model_selection=selection,
                 ),
                 run_context.deps,
             ],
@@ -225,27 +214,27 @@ def _validate_model_request_parameters(self, model_request_parameters: ModelRequ
         if model_request_parameters.allow_image_output:
             raise UserError('Image output is not supported with Temporal because of the 2MB payload size limit.')
 
-    def _current_selection(self) -> ModelSelection:
-        if self._model_selection_var is None:
-            return ModelSelection(model_key='default', model_name=None)
-        selection = self._model_selection_var.get()
-        if selection.model_key is None and selection.model_name is None:
-            return ModelSelection(model_key='default', model_name=None)
-        return selection
+    @contextmanager
+    def using_model(self, selection: str | None) -> Iterator[None]:
+        """Context manager to set the model selection for the duration of a block."""
+        token = self._model_selection_var.set(selection)
+        try:
+            yield
+        finally:
+            self._model_selection_var.reset(token)
+
+    def _current_selection(self) -> str | None:
+        return self._model_selection_var.get()
 
     def _resolve_model(self, params: _RequestParams, deps: Any | None) -> Model:
-        if params.model_key:
-            if params.model_key in self._model_instances:
-                return self._model_instances[params.model_key]
-            raise UserError(
-                f'Model "{params.model_key}" is not registered with this TemporalAgent. '
-                'Register model instances using `additional_models` when constructing the agent.'
-            )
+        selection = params.model_selection
+        if selection is None:
+            return self.wrapped
 
-        if params.model_name:
-            return self._infer_model(params.model_name, params, deps)
+        if selection in self._model_instances:
+            return self._model_instances[selection]
 
-        return self.wrapped
+        return self._infer_model(selection, params, deps)
 
     def _infer_model(self, model_name: str, params: _RequestParams, deps: Any | None) -> Model:
         run_context: RunContext[Any] | None = None