[Temporal - Documentation] Fixing PR errors

Author: scalabreseGD

examples/pydantic_ai_examples/temporal_streaming/agents.py

@@ -3,18 +3,17 @@
 This module defines the agent setup with MCP toolsets, model configuration,
 and custom tools for data analysis.
 """
-
 from datetime import timedelta
 
 from temporalio.common import RetryPolicy
 from temporalio.workflow import ActivityConfig
 
 from pydantic_ai import Agent, FilteredToolset, ModelSettings
+from pydantic_ai.agent import EventStreamHandler
 from pydantic_ai.durable_exec.temporal import TemporalAgent
 from pydantic_ai.mcp import MCPServerStdio
 from pydantic_ai.models.anthropic import AnthropicModel
 from pydantic_ai.providers.anthropic import AnthropicProvider
-
 from .datamodels import AgentDependencies
 
 
@@ -35,43 +34,40 @@ async def get_mcp_toolsets() -> dict[str, FilteredToolset]:
     return {'yahoo': yf_server.filtered(lambda ctx, tool_def: True)}
 
 
-async def get_claude_model(parallel_tool_calls: bool = True, **env_vars):
+async def get_claude_model(parallel_tool_calls: bool = True, **kwargs):
     """
     Create and configure the Claude model.
 
     Args:
         parallel_tool_calls: Whether to enable parallel tool calls.
-        **env_vars: Environment variables including API keys.
+        **kwargs: Environment variables including API keys.
 
     Returns:
         Configured AnthropicModel instance.
     """
-    model_name = 'claude-sonnet-4-5-20250929'
-    api_key = env_vars.get('anthropic_api_key')
-    model = AnthropicModel(
+    model_name: str = 'claude-sonnet-4-5-20250929'
+    api_key: str | None = kwargs.get('anthropic_api_key', None)
+    model: AnthropicModel = AnthropicModel(
         model_name=model_name,
         provider=AnthropicProvider(api_key=api_key),
         settings=ModelSettings(
-            **{
-                'temperature': 0.5,
-                'n': 1,
-                'max_completion_tokens': 64000,
-                'max_tokens': 64000,
-                'parallel_tool_calls': parallel_tool_calls,
-            }
+            temperature=0.5,
+            max_tokens=64000,
+            parallel_tool_calls=parallel_tool_calls,
         ),
     )
 
     return model
 
 
-async def build_agent(stream_handler=None, **env_vars):
+async def build_agent(stream_handler: EventStreamHandler,
+                      **kwargs) -> TemporalAgent:
     """
     Build and configure the agent with tools and temporal settings.
 
     Args:
         stream_handler: Optional event stream handler for streaming responses.
-        **env_vars: Environment variables including API keys.
+        **kwargs: Environment variables including API keys.
 
     Returns:
         TemporalAgent instance ready for use in Temporal workflows.
@@ -84,7 +80,7 @@ async def build_agent(stream_handler=None, **env_vars):
     toolsets = await get_mcp_toolsets()
     agent = Agent(
         name=agent_name,
-        model=await get_claude_model(**env_vars),
+        model=await get_claude_model(**kwargs),
         toolsets=[*toolsets.values()],
         system_prompt=system_prompt,
         event_stream_handler=stream_handler,

examples/pydantic_ai_examples/temporal_streaming/main.py

@@ -7,27 +7,28 @@
 import asyncio
 import os
 import uuid
+from typing import Any, Union
 
-from pydantic_ai.durable_exec.temporal import AgentPlugin, PydanticAIPlugin
-from temporalio.client import Client
+from temporalio.client import Client, WorkflowHandle
 from temporalio.worker import Worker
 
+from pydantic_ai.durable_exec.temporal import AgentPlugin, PydanticAIPlugin
 from .agents import build_agent
 from .datamodels import EventKind, EventStream
 from .streaming_handler import streaming_handler
 from .utils import read_config_yml
 from .workflow import YahooFinanceSearchWorkflow
 
 
-async def poll_events(workflow_handle):
+async def poll_events(workflow_handle: WorkflowHandle[Any, str]) -> None:
     """
     Poll for events from the workflow and print them.
 
     Args:
         workflow_handle: Handle to the running workflow.
     """
     while True:
-        event = await workflow_handle.query('event_stream', result_type=EventStream | None)
+        event = await workflow_handle.query('event_stream', result_type=Union[EventStream | None])
         if event is None:
             await asyncio.sleep(0.1)
             continue
@@ -41,7 +42,7 @@ async def poll_events(workflow_handle):
             print(f'\n--- Event ---\n{event.content}\n')
 
 
-async def main():
+async def main() -> None:
     """
     Main function to set up and run the Temporal workflow.
 
@@ -76,11 +77,12 @@ async def main():
     ):
         # Execute the workflow
         workflow_id = f'yahoo-finance-search-{uuid.uuid4()}'
-        workflow_handle = await client.start_workflow(
-            YahooFinanceSearchWorkflow.run,
-            args=['What are the latest financial metrics for Apple (AAPL)?'],
+        workflow_handle: WorkflowHandle[Any, str] = await client.start_workflow(
+            'YahooFinanceSearchWorkflow',
+            arg='What are the latest financial metrics for Apple (AAPL)?',
             id=workflow_id,
             task_queue=task_queue,
+            result_type=str
         )
 
         print(f'Started workflow with ID: {workflow_id}')

examples/pydantic_ai_examples/temporal_streaming/streaming_handler.py

@@ -6,6 +6,9 @@
 
 from collections.abc import AsyncIterable
 
+from temporalio import activity
+from temporalio.client import WorkflowHandle
+
 from pydantic_ai import (
     AgentStreamEvent,
     FunctionToolCallEvent,
@@ -15,17 +18,15 @@
     TextPart,
     TextPartDelta,
     ThinkingPartDelta,
-    ToolCallPart,
+    ToolCallPart, RunContext,
 )
-from temporalio import activity
-
 from .datamodels import AgentDependencies, EventKind, EventStream
 
 
 async def streaming_handler(
-    ctx,
-    event_stream_events: AsyncIterable[AgentStreamEvent],
-):
+        ctx: RunContext,
+        event_stream_events: AsyncIterable[AgentStreamEvent],
+) -> None:
     """
     Handle streaming events from the agent.
 
@@ -37,8 +38,11 @@ async def streaming_handler(
         ctx: The run context containing dependencies.
         event_stream_events: Async iterable of agent stream events.
     """
-    output = ''
-    output_tool_delta = dict(
+    if not activity.in_activity():
+        return
+
+    output: str = ''
+    output_tool_delta: dict[str, str] = dict(
         tool_call_id='',
         tool_name_delta='',
         args_delta='',
@@ -78,19 +82,10 @@ async def streaming_handler(
         output += f'\nTool Name: {output_tool_delta["tool_name_delta"]}'
         output += f'\nTool Args: {output_tool_delta["args_delta"]}'
 
-    events = []
-
-    # Create event stream if there's output
-    if output:
-        event = EventStream(kind=EventKind.EVENT, content=output)
-        events.append(event)
-
     # Send events to workflow if running in an activity
-    if activity.in_activity():
-        deps: AgentDependencies = ctx.deps
+    deps: AgentDependencies = ctx.deps
 
-        workflow_id = deps.workflow_id
-        run_id = deps.run_id
-        workflow_handle = activity.client().get_workflow_handle(workflow_id=workflow_id, run_id=run_id)
-        for event in events:
-            await workflow_handle.signal('append_event', arg=event)
+    workflow_id: str = deps.workflow_id
+    run_id: str = deps.run_id
+    workflow_handle: WorkflowHandle = activity.client().get_workflow_handle(workflow_id=workflow_id, run_id=run_id)
+    await workflow_handle.signal('append_event', arg=EventStream(kind=EventKind.EVENT, content=output))

examples/pydantic_ai_examples/temporal_streaming/utils.py

@@ -2,11 +2,12 @@
 
 import os
 from copy import copy
+from typing import Any
 
 import yaml
 
 
-def recursively_modify_api_key(conf):
+def recursively_modify_api_key(conf) -> dict[str, Any]:
     """
     Recursively replace API key placeholders with environment variable values.
 
@@ -41,7 +42,7 @@ def inner(_conf):
     return copy_conf
 
 
-def read_config_yml(path):
+def read_config_yml(path) -> dict[str, Any]:
     """
     Read and process a YAML configuration file.
 

examples/pydantic_ai_examples/temporal_streaming/workflow.py

@@ -10,9 +10,9 @@
 from datetime import timedelta
 from typing import Any
 
-from pydantic_ai import UsageLimits
 from temporalio import activity, workflow
 
+from pydantic_ai import UsageLimits
 from .agents import build_agent
 from .datamodels import AgentDependencies, EventKind, EventStream
 from .streaming_handler import streaming_handler
@@ -33,7 +33,7 @@ def __init__(self):
         self.events: deque[EventStream] = deque()
 
     @workflow.run
-    async def run(self, user_prompt: str):
+    async def run(self, user_prompt: str) -> str:
         """
         Execute the agent with the given user prompt.
 
@@ -78,7 +78,7 @@ async def run(self, user_prompt: str):
 
     @staticmethod
     @activity.defn(name='retrieve_env_vars')
-    async def retrieve_env_vars():
+    async def retrieve_env_vars() -> dict[str, Any]:
         """
         Retrieve environment variables from configuration file.