AWS Lambda durable functionsがPydantic AIと統合されたので試してみる

Lambda ランタイム: python3.13 (arm64)
pydantic-ai-harness 0.36.0
pydantic-ai-slim 2.51.0
aws-durable-execution-sdk-python 2.0.1
AWS SAM CLI 1.166.2
AWS CLI 2.37.7
リージョン: ap-northeast-1
モデル: apac.amazon.nova-lite-v1:0 (Amazon Bedrock)

import os
import time
from typing import Any

from aws_durable_execution_sdk_python import DurableContext, durable_execution
from pydantic_ai import Agent
from pydantic_ai.usage import UsageLimits
from pydantic_ai_harness.aws_lambda import AWSLambdaDurability, durable_agent_handler

MODEL_ID = os.environ.get("MODEL_ID", "apac.amazon.nova-lite-v1:0")

agent = Agent(
    f"bedrock:{MODEL_ID}",
    name="travel",
    instructions=(
        "あなたは旅行アシスタントです。get_weather と get_exchange_rate で取得した情報を使って、"
        "指定された都市への旅行アドバイスを日本語で3文以内で答えてください。"
        "通貨コードは ISO 4217（例: THB）で指定してください。思考過程は出力しないでください。"
    ),
    capabilities=[AWSLambdaDurability()],
)

@agent.tool_plain
def get_weather(city: str) -> str:
    """指定した都市の今日の天気を返す（外部 API の代わりのモック）。"""
    delay = float(os.environ.get("WEATHER_API_DELAY", "0"))
    time.sleep(delay)  # 外部 API の応答待ちを模擬する
    return f"{city} の今日の天気は晴れ、最高気温は 28 度です。"

@agent.tool_plain
def get_exchange_rate(currency: str) -> str:
    """指定した通貨コード（ISO 4217）の対円レートを返す（外部 API の代わりのモック）。"""
    rates = {"USD": 150.2, "EUR": 162.8, "THB": 4.3}
    code = currency.strip().upper()
    rate = rates.get(code)
    if rate is None:
        return f"{code} のレートは取得できませんでした。"
    return f"1 {code} = {rate} 円です。"

@durable_execution
@durable_agent_handler
async def handler(event: dict[str, Any], context: DurableContext) -> dict[str, Any]:
    result = await agent.run(
        str(event["prompt"]),
        usage_limits=UsageLimits(request_limit=6),
    )
    return {
        "output": result.output,
        "requests": result.usage.requests,
        "tool_calls": result.usage.tool_calls,
    }

pydantic-ai-harness[aws-lambda]==0.36.0
pydantic-ai-slim[bedrock]==2.51.0
aws-durable-execution-sdk-python==2.0.1

AWSTemplateFormatVersion: "2010-09-09"
Transform: AWS::Serverless-2016-10-31
Description: AWS Lambda durable functions x Pydantic AI sample

Parameters:
  ModelId:
    Type: String
    Default: apac.amazon.nova-lite-v1:0
  FunctionTimeout:
    Type: Number
    Default: 60
  WeatherApiDelay:
    Type: String
    Default: "0"

Resources:
  DurableAgentFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: durable-pydantic-ai-agent
      CodeUri: src/
      Handler: handler.handler
      Runtime: python3.13
      Architectures:
        - arm64
      MemorySize: 1024
      Timeout: !Ref FunctionTimeout
      # DurableConfig を書くと、SAM が生成する実行ロールに AWSLambdaBasicDurableExecutionRolePolicy が付く
      DurableConfig:
        ExecutionTimeout: 900
        RetentionPeriodInDays: 7
      Policies:
        - Statement:
            - Effect: Allow
              Action:
                - bedrock:InvokeModel
                - bedrock:InvokeModelWithResponseStream
              Resource: "*"
      Environment:
        Variables:
          MODEL_ID: !Ref ModelId
          WEATHER_API_DELAY: !Ref WeatherApiDelay

Outputs:
  FunctionName:
    Value: !Ref DurableAgentFunction
  FunctionArn:
    Value: !GetAtt DurableAgentFunction.Arn

$ sam build
$ sam deploy --stack-name durable-pydantic-ai --region ap-northeast-1 \
    --resolve-s3 --capabilities CAPABILITY_IAM --no-confirm-changeset

$ aws lambda invoke --function-name 'durable-pydantic-ai-agent:$LATEST' \
    --region ap-northeast-1 --cli-binary-format raw-in-base64-out \
    --payload '{"prompt": "来週バンコクに行きます。天気とタイの通貨レートを調べて助言してください。"}' \
    response.json
{
    "StatusCode": 200,
    "ExecutedVersion": "$LATEST",
    "DurableExecutionArn": "arn:aws:lambda:ap-northeast-1:xxxxxxxxxxxx:function:durable-pydantic-ai-agent:$LATEST/durable-execution/892c83be-0f87-41f2-815d-a837400db650/10d026e5-95bf-3e4c-a6e3-3cddf200f369"
}

{
  "output": "来週のバンコクの天気は晴れで、最高気温は28度です。1タイバーツは4.3円に換算されます。",
  "requests": 2,
  "tool_calls": 2
}

$ aws lambda get-durable-execution-history --region ap-northeast-1 \
    --durable-execution-arn 'arn:aws:lambda:ap-northeast-1:xxxxxxxxxxxx:function:durable-pydantic-ai-agent:$LATEST/durable-execution/892c83be-0f87-41f2-815d-a837400db650/10d026e5-95bf-3e4c-a6e3-3cddf200f369'

$ aws lambda get-durable-execution-history --region ap-northeast-1 \
    --durable-execution-arn '（省略）' --include-execution-data \
    --query 'Events[?Name==`travel__model.request` && EventType==`StepSucceeded`] | [0].StepSucceededDetails.Result.Payload' --output text | head -c 300
{"t":"m","v":{"parts":{"t":"l","v":[{"t":"m","v":{"content":{"t":"s","v":"<thinking>\u6765\u9031\u30d0\u30f3\u30b3\u30af\u3078\u306e\u65c5\u884c\u306b\u5929\u6c17\u3068\u30bf\u30a4\u306e\u901a\u8ca8\u30ec\u30fc\u30c8\u306e\u60c5\u5831\u304c\u5fc5\u8981\u3067\u3059\u3002get_weather\u3068get_exchange_

$ sam deploy --stack-name durable-pydantic-ai --region ap-northeast-1 \
    --resolve-s3 --capabilities CAPABILITY_IAM --no-confirm-changeset \
    --parameter-overrides FunctionTimeout=20 WeatherApiDelay=30

$ aws lambda invoke --function-name 'durable-pydantic-ai-agent:$LATEST' \
    --region ap-northeast-1 --cli-binary-format raw-in-base64-out \
    --payload '{"prompt": "来週バンコクに行きます。天気とタイの通貨レートを調べて助言してください。"}' \
    response.json &
$ sleep 28
$ aws lambda update-function-configuration --function-name durable-pydantic-ai-agent \
    --region ap-northeast-1 \
    --environment "Variables={MODEL_ID=apac.amazon.nova-lite-v1:0,WEATHER_API_DELAY=0}"

{
  "output": "バンコクの天気は晴れで最高気温が28度です。1 THBは約4.3円です。",
  "requests": 2,
  "tool_calls": 2
}

"Error": {
  "Payload": {
    "ErrorMessage": "RequestId: 6e30ff2f-5c59-4201-9999-21320f9e794f Error: Task timed out after 20.00 seconds",
    "ErrorType": "Sandbox.Timedout"
  },
  "Truncated": false
}

$ sam delete --stack-name durable-pydantic-ai --region ap-northeast-1 --no-prompts