diff --git a/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json b/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json index 86878529..2b78646f 100644 --- a/packages/aws-durable-execution-sdk-python-examples/examples-catalog.json +++ b/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" } ] } diff --git a/packages/aws-durable-execution-sdk-python-examples/pyproject.toml b/packages/aws-durable-execution-sdk-python-examples/pyproject.toml index 99cb61bf..333a63d3 100644 --- a/packages/aws-durable-execution-sdk-python-examples/pyproject.toml +++ b/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"] diff --git a/packages/aws-durable-execution-sdk-python-examples/src/pydantic_ai_agent/lambda_durability.py b/packages/aws-durable-execution-sdk-python-examples/src/pydantic_ai_agent/lambda_durability.py new file mode 100644 index 00000000..2e954c91 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-examples/src/pydantic_ai_agent/lambda_durability.py @@ -0,0 +1,209 @@ +# 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(): + try: + task = loop.create_task(coro_factory(), context=step_ctx) + task.add_done_callback(on_done) + except BaseException as e: + future.set_exception(e) + + 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 diff --git a/packages/aws-durable-execution-sdk-python-examples/src/pydantic_ai_agent/pydantic-ai-agent.zip b/packages/aws-durable-execution-sdk-python-examples/src/pydantic_ai_agent/pydantic-ai-agent.zip new file mode 100644 index 00000000..d7123f5c Binary files /dev/null and b/packages/aws-durable-execution-sdk-python-examples/src/pydantic_ai_agent/pydantic-ai-agent.zip differ diff --git a/packages/aws-durable-execution-sdk-python-examples/src/pydantic_ai_agent/pydantic_ai_agent.py b/packages/aws-durable-execution-sdk-python-examples/src/pydantic_ai_agent/pydantic_ai_agent.py new file mode 100644 index 00000000..a96bc9cd --- /dev/null +++ b/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 diff --git a/packages/aws-durable-execution-sdk-python-examples/test/pydantic_ai_agent/test_pydantic_ai_agent.py b/packages/aws-durable-execution-sdk-python-examples/test/pydantic_ai_agent/test_pydantic_ai_agent.py new file mode 100644 index 00000000..33236da3 --- /dev/null +++ b/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)