Fix using sync history processors, instructions functions, and output functions with TemporalAgent #3704
Conversation
…oralAgent context management
Docs Preview
|
| with super().override(model=self._model, toolsets=self._toolsets, tools=[]): | ||
| token = self._temporal_overrides_active.set(True) | ||
| temporal_active_token = self._temporal_overrides_active.set(True) | ||
| blocking_token = _prefer_blocking_execution.set(True) |
There was a problem hiding this comment.
Let's make this a contextmanager exposed by the utils module, so that we don't deal with the contextvar directly here. See the current_run_context feature in #3537
There was a problem hiding this comment.
Then we can add that contextmanager to the with statement above
tests/test_temporal.py Outdated
| return messages[1:] if len(messages) > 1 else messages | ||
| | ||
| | ||
| agent_with_sync_history_processor = Agent( |
There was a problem hiding this comment.
Move this to just ahead of the workflow that uses it, to keep the test file consistently ordered
tests/test_temporal.py Outdated
| simple_temporal_agent = TemporalAgent(simple_agent, activity_config=BASE_ACTIVITY_CONFIG) | ||
| | ||
| | ||
| def drop_first_message_sync(messages: list[ModelMessage]) -> list[ModelMessage]: |
There was a problem hiding this comment.
Let's make this an inline lambda as it's very simple
There was a problem hiding this comment.
pyright always complains about lambdas, and the length causes a line break, so I'd prefer to keep the extra 1-2 lines with proper types and a nice name
tests/test_utils.py Outdated
| | ||
| | ||
| async def test_run_in_executor_with_blocking_execution_enabled() -> None: | ||
| from pydantic_ai._utils import _prefer_blocking_execution # pyright: ignore[reportPrivateUsage] |
There was a problem hiding this comment.
Same as above; let's not use the private var, but a new public contextmanager
| _P = ParamSpec('_P') | ||
| _R = TypeVar('_R') | ||
| | ||
| _prefer_blocking_execution: ContextVar[bool] = ContextVar('_prefer_blocking_execution', default=False) |
There was a problem hiding this comment.
"prefer" makes it sound like this expresses a preference, that may or may not be respected. So "enable" or no prefix at all would be better
| | ||
| | ||
| @contextmanager | ||
| def blocking_execution() -> Iterator[None]: |
There was a problem hiding this comment.
with _utils.blocking_execution() makes it seems like all execution will be blocking, but that's not true because of async etc. So maybe we can rename this to something like with disable_threads()?
| | ||
| @contextmanager | ||
| def _temporal_overrides(self) -> Iterator[None]: | ||
| from pydantic_ai._utils import blocking_execution |
There was a problem hiding this comment.
To the top please! Imports inside methods are exclusively for circular import issues, and files that depend on optional packages
DouweM changed the title TemporalAgent 
Problem
Sync history processors fail inside Temporal workflows because
anyio.to_thread.run_synctries to create threads, which Temporal's sandboxed event loop doesn't support. This causesNotImplementedError.Solution
Use a ContextVar to control when
run_in_executor()should execute sync functions directly (blocking) vs using threading. This follows the approach discussed in the issue comments and approved by @DouweM.Changes
_utils.py: Added_prefer_blocking_executionContextVar, removed Temporal detection codetemporal/_agent.py: Set ContextVar in_temporal_overrides()context managertest_utils.py: Updated test to use ContextVar instead of mocking Temporal internalstest_temporal.py: Added integration test for sync history processors in workflowsTesting