Add Pydantic AI Durable Capability Example

packages/aws-durable-execution-sdk-python-examples/examples-catalog.json

@@ -763,6 +763,20 @@
         "FILESYSTEM_MOUNT_PATH": "/mnt/s3"
       },
       "path": "./src/filesystem_serdes/filesystem_serdes_preview.py"
+    },
+    {
+      "name": "Pydantic AI Durable Agent",
+      "description": "Checkpoints Pydantic AI model requests and tool calls as durable steps",
+      "handler": "pydantic_ai_agent.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "environment": {
+        "PYDANTIC_AI_MODEL": "us.amazon.nova-micro-v1:0"
+      },
+      "path": "./src/pydantic_ai_agent/pydantic_ai_agent.py"
     }
   ]
 }

packages/aws-durable-execution-sdk-python-examples/pyproject.toml

@@ -13,6 +13,11 @@ dependencies = [
   "aws-durable-execution-sdk-python-otel"
 ]
 
+[project.optional-dependencies]
+pydantic-ai = [
+  "pydantic-ai-slim[bedrock]",
+]
+
 [tool.hatch.build.targets.wheel]
 packages = ["src"]
 

packages/aws-durable-execution-sdk-python-examples/src/pydantic_ai_agent/lambda_durability.py

@@ -0,0 +1,206 @@
+# The capability bridges the sync Lambda durable functions API with the async
+# agent, using two threads connected by a queue:
+#
+#   - Main thread: runs the durable handler and its steps.
+#   - Background thread: runs the async agent and its model/tool calls.
+#
+# On each model/tool call the agent enqueues a step request. The main thread
+# consumes the queue, runs the durable step, and dispatches the actual async
+# call back to the background loop.
+
+import asyncio
+import contextvars
+import threading
+from concurrent.futures import Future
+from queue import Queue
+from typing import Any
+
+from aws_durable_execution_sdk_python import DurableContext
+from aws_durable_execution_sdk_python.config import StepConfig
+
+try:
+    from pydantic_ai.capabilities import AbstractCapability
+except ImportError:
+
+    class AbstractCapability:  # type: ignore[no-redef]
+        pass
+
+
+# Durable steps checkpoint JSON data by default.
+# Use these to convert model responses and retry prompts from/to JSON.
+def _model_response_adapter() -> Any:
+    from pydantic import TypeAdapter
+    from pydantic_ai.messages import ModelResponse
+
+    return TypeAdapter(ModelResponse)
+
+
+def _retry_prompt_adapter() -> Any:
+    from pydantic import TypeAdapter
+    from pydantic_ai.messages import RetryPromptPart
+
+    return TypeAdapter(RetryPromptPart)
+
+
+# Persistent background-thread event loop that runs the agent and its
+# model/tool calls, reused across warm invocations so loop-bound async
+# resources (e.g. a cached HTTP client) stay valid.
+class _AgentLoop:
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._loop: asyncio.AbstractEventLoop | None = None
+
+    def get(self) -> asyncio.AbstractEventLoop:
+        with self._lock:
+            if self._loop is None or self._loop.is_closed():
+                self._loop = asyncio.new_event_loop()
+                threading.Thread(target=self._loop.run_forever, daemon=True).start()
+            return self._loop
+
+
+_agent_loop = _AgentLoop()
+
+
+# Queue-based bridge: the agent runs on the background thread and sends step
+# requests to the queue. The main handler thread consumes the queue and runs
+# context.step(), so every durable step is created on the same thread as the
+# handler (one continuous context).
+class _StepQueue:
+    def __init__(self):
+        self._queue: Queue = Queue()
+
+    # Run on background thread: add step to queue, await result from main thread
+    async def run_step(self, step_body, name, config, context):
+        loop = asyncio.get_running_loop()
+        future = loop.create_future()
+
+        def reply(success, value):
+            fn = future.set_result if success else future.set_exception
+            loop.call_soon_threadsafe(fn, value)
+
+        # Copy the agent's contextvars to the durable step.
+        ctx = contextvars.copy_context()
+        self._queue.put(("step", step_body, name, config, context, reply, ctx))
+        return await future
+
+    def finish(self, result=None, error=None):
+        self._queue.put(("done", result, error))
+
+    # Run on main thread: execute queued durable steps until finish()
+    def consume(self):
+        while True:
+            item = self._queue.get()
+            if item[0] == "done":
+                _, result, error = item
+                if error is not None:
+                    raise error
+                return result
+            _, step_body, name, config, context, reply, ctx = item
+            try:
+                reply(True, ctx.run(context.step, step_body, name=name, config=config))
+            except BaseException as e:
+                reply(False, e)
+
+
+class LambdaDurability(AbstractCapability):
+    def __init__(
+        self,
+        context: DurableContext,
+        step_config: StepConfig | None = None,
+        tool_configs: dict[str, StepConfig] | None = None,
+    ):
+        self.context = context
+        self.step_config = step_config
+        self.tool_configs = tool_configs or {}
+        self.step_queue: _StepQueue | None = None
+
+    async def wrap_run(self, ctx, *, handler):
+        from pydantic_ai.tool_manager import ToolManager
+
+        # Run tool calls one at a time. Durable step IDs are assigned by
+        # call order, so concurrent step() calls could race and break replay.
+        with ToolManager.parallel_execution_mode("sequential"):
+            return await handler()
+
+    async def _durable_step(self, coro_factory, name, config):
+        loop = asyncio.get_running_loop()
+
+        # step_body runs on the main thread. It schedules the async model/tool
+        # call back onto the background thread (which stays free while we await).
+        def step_body(_step_ctx):
+            future: Future = Future()
+
+            # Copy the durable step's contextvars to model/tool calls.
+            step_ctx = contextvars.copy_context()
+
+            def on_done(task):
+                if task.exception():
+                    future.set_exception(task.exception())
+                else:
+                    future.set_result(task.result())
+
+            def schedule():
+                task = loop.create_task(coro_factory(), context=step_ctx)
+                task.add_done_callback(on_done)
+
+            loop.call_soon_threadsafe(schedule)
+            return future.result()
+
+        return await self.step_queue.run_step(step_body, name, config, self.context)
+
+    async def wrap_model_request(self, ctx, *, request_context, handler):
+        adapter = _model_response_adapter()
+
+        async def call_model():
+            response = await handler(request_context)
+            return adapter.dump_python(response, mode="json")
+
+        checkpointed = await self._durable_step(
+            call_model, "model.request", self.step_config
+        )
+        return adapter.validate_python(checkpointed)
+
+    async def wrap_tool_execute(self, ctx, *, call, tool_def, args, handler):
+        from pydantic_ai.exceptions import ToolRetryError
+
+        adapter = _retry_prompt_adapter()
+
+        # A ToolRetryError is a control signal, not a failure. Checkpoint it as
+        # a normal result inside the step, then re-raise it outside so the step
+        # is not recorded as failed.
+        async def call_tool():
+            try:
+                return {"retry": False, "value": await handler(args)}
+            except ToolRetryError as e:
+                retry_prompt = adapter.dump_python(e.tool_retry, mode="json")
+                return {"retry": True, "retry_prompt": retry_prompt}
+
+        config = self.tool_configs.get(tool_def.name, self.step_config)
+        result = await self._durable_step(call_tool, f"tool.{tool_def.name}", config)
+        if result["retry"]:
+            raise ToolRetryError(adapter.validate_python(result["retry_prompt"]))
+        return result["value"]
+
+
+# Run the agent on the background thread while the main thread consumes
+# and executes its durable steps from the queue.
+def run_durable(coro_factory, capability: LambdaDurability):
+    step_queue = _StepQueue()
+    capability.step_queue = step_queue
+    loop = _agent_loop.get()
+
+    async def run_agent():
+        try:
+            step_queue.finish(result=await coro_factory())
+        except BaseException as e:
+            step_queue.finish(error=e)
+
+    # Copy the handler's contextvars to the agent. Fresh copy each invocation
+    # so warm starts do not reuse a stale context.
+    ctx = contextvars.copy_context()
+
+    def schedule():
+        loop.create_task(run_agent(), context=ctx)
+
+    loop.call_soon_threadsafe(schedule)
+    return step_queue.consume()   # blocks on the main thread until the agent finishes

packages/aws-durable-execution-sdk-python-examples/src/pydantic_ai_agent/pydantic_ai_agent.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import os
+
+from aws_durable_execution_sdk_python import DurableContext, durable_execution
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.models.bedrock import BedrockConverseModel
+
+try:
+    from lambda_durability import LambdaDurability, run_durable
+except ImportError:
+    from src.pydantic_ai_agent.lambda_durability import LambdaDurability, run_durable
+
+MODEL_ID = os.environ.get("PYDANTIC_AI_MODEL", "us.amazon.nova-micro-v1:0")
+
+agent = Agent(BedrockConverseModel(MODEL_ID))
+
+
+@agent.tool
+def get_weather(ctx: RunContext, city: str) -> str:
+    return f"It is sunny and 22C in {city}."
+
+
+@durable_execution
+def handler(event: dict, context: DurableContext) -> str:
+    prompt = event.get("prompt", "What is the weather in Vancouver?")
+    capability = LambdaDurability(context)
+    result = run_durable(
+        lambda: agent.run(prompt, capabilities=[capability]),
+        capability,
+    )
+    return result.output

packages/aws-durable-execution-sdk-python-examples/test/pydantic_ai_agent/test_pydantic_ai_agent.py

@@ -0,0 +1,84 @@
+"""Tests for the Pydantic AI durable agent example."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+
+import pytest
+
+from aws_durable_execution_sdk_python.execution import InvocationStatus
+from src.pydantic_ai_agent import lambda_durability
+from test.conftest import deserialize_operation_payload
+
+
+class FakeStepContext:
+    pass
+
+
+class FakeDurableContext:
+    """Records the durable steps run against it and the thread they run on."""
+
+    def __init__(self) -> None:
+        self.steps: list[str] = []
+        self.step_thread_ids: list[int] = []
+
+    def step(self, func, name=None, config=None):
+        self.step_thread_ids.append(threading.get_ident())
+        self.steps.append(name)
+        return func(FakeStepContext())
+
+
+@pytest.mark.example
+def test_run_durable_runs_steps_on_the_handler_thread():
+    context = FakeDurableContext()
+    handler_thread_id = threading.get_ident()
+    agent_thread_ids: list[int] = []
+
+    capability = lambda_durability.LambdaDurability(context)
+
+    async def call_model():
+        agent_thread_ids.append(threading.get_ident())
+        return {"answer": 42}
+
+    async def run_agent():
+        # Exercise a single durable step through the queue bridge.
+        return await capability._durable_step(call_model, "model.request", None)
+
+    result = lambda_durability.run_durable(run_agent, capability)
+
+    assert result == {"answer": 42}
+    # The durable step ran on the handler (main) thread.
+    assert context.step_thread_ids == [handler_thread_id]
+    assert context.steps == ["model.request"]
+    # The model call ran on the background agent thread, not the handler thread.
+    assert agent_thread_ids and agent_thread_ids[0] != handler_thread_id
+
+
+@pytest.mark.example
+def test_agent_loop_thread_is_reused_across_invocations():
+    def loop_identity() -> int:
+        return id(lambda_durability._agent_loop.get())
+
+    assert loop_identity() == loop_identity()
+
+
+@pytest.mark.example
+@pytest.mark.durable_execution(
+    lambda_function_name="Pydantic AI Durable Agent",
+)
+def test_pydantic_ai_agent_runs_in_cloud(durable_runner):
+    if durable_runner.mode != "cloud":
+        pytest.skip("Pydantic AI Bedrock example only runs in cloud mode")
+
+    with durable_runner:
+        result = durable_runner.run(
+            input={"prompt": "What is the weather in Vancouver?"},
+            timeout=120,
+        )
+
+    assert result.status is InvocationStatus.SUCCEEDED
+    assert deserialize_operation_payload(result.result)
+    operation_names = {operation.name for operation in result.operations}
+    assert any(name and name.startswith("model.request") for name in operation_names)
+    assert any(name and name.startswith("tool.") for name in operation_names)