Add AI Workflow Lambda durable function snippets

lambda-durable-ai-agent-with-tools/example.py

@@ -0,0 +1,92 @@
+from typing import Any, Callable
+
+import boto3
+from aws_durable_execution_sdk_python import DurableContext, durable_execution
+
+MODEL_ID = "us.amazon.nova-pro-v1:0"
+bedrock = boto3.client("bedrock-runtime")
+
+
+class AgentTool:
+    def __init__(self, tool_spec: dict, execute: Callable[[dict, DurableContext], str]):
+        self.tool_spec = tool_spec
+        self.execute = execute
+
+
+TOOLS: list[AgentTool] = [
+    AgentTool(
+        tool_spec={
+            "name": "get_weather",
+            "description": "Get the current weather for a location.",
+            "inputSchema": {
+                "json": {
+                    "type": "object",
+                    "properties": {"location": {"type": "string"}},
+                    "required": ["location"],
+                }
+            },
+        },
+        execute=lambda input, ctx: f"The weather in {input.get('location', 'unknown')} is sunny, 72°F.",
+    ),
+    AgentTool(
+        tool_spec={
+            "name": "wait_for_human_review",
+            "description": "Request human review and wait for response.",
+            "inputSchema": {
+                "json": {
+                    "type": "object",
+                    "properties": {"question": {"type": "string"}},
+                    "required": ["question"],
+                }
+            },
+        },
+        execute=lambda input, ctx: ctx.wait_for_callback(
+            lambda callback_id, _: print(f"Review needed: {input.get('question')}"),
+            "human_review",
+        ),
+    ),
+]
+
+
+@durable_execution
+def handler(event: dict, context: DurableContext):
+    prompt = event.get("prompt", "What's the weather in Seattle?")
+    messages: list[Any] = [{"role": "user", "content": [{"text": prompt}]}]
+    tools_by_name = {t.tool_spec["name"]: t for t in TOOLS}
+
+    while True:
+        response = context.step(
+            lambda _: bedrock.converse(
+                modelId=MODEL_ID,
+                messages=messages,
+                toolConfig={"tools": [{"toolSpec": t.tool_spec} for t in TOOLS]},
+            ),
+            "converse",
+        )
+
+        output = response.get("output", {}).get("message", {})
+        messages.append(output)
+
+        if response.get("stopReason") == "end_turn":
+            for block in output.get("content", []):
+                if "text" in block:
+                    return block["text"]
+            return ""
+
+        tool_results = []
+        for block in output.get("content", []):
+            if "toolUse" in block:
+                tool_use = block["toolUse"]
+                tool = tools_by_name[tool_use["name"]]
+                result = context.run_in_child_context(
+                    lambda child_ctx: tool.execute(tool_use.get("input", {}), child_ctx),
+                    f"tool:{tool_use['name']}",
+                )
+                tool_results.append({
+                    "toolResult": {
+                        "toolUseId": tool_use["toolUseId"],
+                        "content": [{"text": result}],
+                    }
+                })
+
+        messages.append({"role": "user", "content": tool_results})

lambda-durable-ai-agent-with-tools/example.ts

@@ -0,0 +1,95 @@
+import {
+  type DurableContext,
+  withDurableExecution,
+} from "@aws/durable-execution-sdk-js";
+import {
+  BedrockRuntimeClient,
+  type ContentBlock,
+  ConverseCommand,
+  type Message,
+  type Tool,
+} from "@aws-sdk/client-bedrock-runtime";
+
+const MODEL_ID = "us.amazon.nova-pro-v1:0";
+const bedrock = new BedrockRuntimeClient({});
+
+type AgentTool = {
+  toolSpec: NonNullable<Tool["toolSpec"]>;
+  execute: (input: Record<string, string>, context: DurableContext) => Promise<string>;
+};
+
+const tools: AgentTool[] = [
+  {
+    toolSpec: {
+      name: "get_weather",
+      description: "Get the current weather for a location.",
+      inputSchema: {
+        json: {
+          type: "object",
+          properties: { location: { type: "string" } },
+          required: ["location"],
+        },
+      },
+    },
+    execute: async (input) => `The weather in ${input.location} is sunny, 72°F.`,
+  },
+  {
+    toolSpec: {
+      name: "wait_for_human_review",
+      description: "Request human review and wait for response.",
+      inputSchema: {
+        json: {
+          type: "object",
+          properties: { question: { type: "string" } },
+          required: ["question"],
+        },
+      },
+    },
+    execute: async (input, context) =>
+      context.waitForCallback<string>("human_review", async (callbackId) => {
+        console.log(`Review needed: ${input.question}`);
+      }),
+  },
+];
+
+export const handler = withDurableExecution(
+  async (event: { prompt?: string }, context: DurableContext) => {
+    const prompt = event.prompt ?? "What's the weather in Seattle?";
+    const messages: Message[] = [{ role: "user", content: [{ text: prompt }] }];
+    const toolsByName = Object.fromEntries(tools.map((t) => [t.toolSpec.name, t]));
+
+    while (true) {
+      const response = await context.step("converse", async () => {
+        return bedrock.send(
+          new ConverseCommand({
+            modelId: MODEL_ID,
+            messages,
+            toolConfig: { tools: tools.map((t) => ({ toolSpec: t.toolSpec })) },
+          })
+        );
+      });
+
+      const output = response.output!.message!;
+      messages.push(output);
+
+      if (response.stopReason === "end_turn") {
+        const textBlock = output.content?.find((b): b is ContentBlock.TextMember => "text" in b);
+        return textBlock?.text ?? "";
+      }
+
+      const toolResults: ContentBlock[] = [];
+      for (const block of output.content ?? []) {
+        if ("toolUse" in block && block.toolUse) {
+          const { toolUseId, name, input } = block.toolUse;
+          const tool = toolsByName[name!];
+          const result = await context.runInChildContext(`tool:${name}`, async (childContext) => {
+            return tool.execute(input as Record<string, string>, childContext);
+          });
+          toolResults.push({ toolResult: { toolUseId, content: [{ text: result }] } });
+        }
+      }
+
+      messages.push({ role: "user", content: toolResults });
+    }
+  }
+);

lambda-durable-ai-agent-with-tools/snippet-data.json

@@ -0,0 +1,66 @@
+{
+  "title": "Agent with Tools using AWS Lambda durable functions",
+  "description": "Agentic loop where the LLM can call tools, including tools that suspend for human input, using AWS Lambda durable functions",
+  "type": "Integration",
+  "services": ["lambda", "bedrock"],
+  "languages": ["Python", "TypeScript"],
+  "tags": ["ai", "durable-functions", "bedrock", "genai", "agent", "tools", "function-calling"],
+  "introBox": {
+    "headline": "How it works",
+    "text": [
+      "This pattern demonstrates an agentic loop with AWS Lambda durable functions. The LLM can call tools, and each model call and tool execution is checkpointed for deterministic replay.",
+      "Tools can use durable context features like waitForCallback to suspend the agent while awaiting external input.",
+      "Benefits: Complex agent loops are expressed as simple sequential code. Each tool call is checkpointed for resilient execution. Tools can suspend for human input without consuming compute resources."
+    ]
+  },
+  "gitHub": {
+    "template": {
+      "repoURL": "https://github.com/aws-samples/sample-ai-workflows-in-aws-lambda-durable-functions"
+    }
+  },
+  "snippets": [
+    {
+      "title": "Runtimes",
+      "codeTabs": [
+        {
+          "id": "Python",
+          "title": "Usage Example with Python:",
+          "description": "Agent with tools using Lambda durable functions in Python.",
+          "snippets": [
+            {
+              "snippetPath": "example.py",
+              "language": "py"
+            }
+          ]
+        },
+        {
+          "id": "TypeScript",
+          "title": "Usage Example with TypeScript:",
+          "description": "Agent with tools using Lambda durable functions in TypeScript.",
+          "snippets": [
+            {
+              "snippetPath": "example.ts",
+              "language": "ts"
+            }
+          ]
+        }
+      ]
+    }
+  ],
+  "resources": {
+    "bullets": [
+      {
+        "text": "AWS Lambda durable functions",
+        "link": "https://aws.amazon.com/lambda/lambda-durable-functions/"
+      }
+    ]
+  },
+  "authors": [
+    {
+      "headline": "Presented by Connor Kirkpatrick",
+      "name": "Connor Kirkpatrick",
+      "bio": "Solution Engineer at AWS",
+      "linkedin": "connorkirkpatrick"
+    }
+  ]
+}

lambda-durable-ai-human-review/example.py

@@ -0,0 +1,57 @@
+import json
+
+import boto3
+from aws_durable_execution_sdk_python import DurableContext, durable_execution
+from aws_durable_execution_sdk_python.config import Duration, WaitForCallbackConfig
+from pydantic import BaseModel
+
+MODEL_ID = "us.amazon.nova-pro-v1:0"
+bedrock = boto3.client("bedrock-runtime")
+
+
+def converse(model_id: str, prompt: str) -> str:
+    response = bedrock.converse(
+        modelId=model_id,
+        messages=[{"role": "user", "content": [{"text": prompt}]}],
+    )
+    return response["output"]["message"]["content"][0]["text"]
+
+
+class ReviewResult(BaseModel):
+    approved: bool
+    notes: str | None = None
+
+
+def send_for_review(callback_id: str, document: str, extracted_fields: str):
+    print(f"Review needed for document. Callback ID: {callback_id}")
+    print(f"Extracted fields: {extracted_fields}")
+
+
+@durable_execution
+def handler(event: dict, context: DurableContext):
+    document = event.get("document", "Sample invoice with amount $1,234.56")
+
+    extracted_fields = context.step(
+        lambda _: converse(
+            MODEL_ID,
+            f'Extract key fields from this document as JSON: "{document}"',
+        ),
+        "extract fields",
+    )
+
+    review_result_str = context.wait_for_callback(
+        lambda callback_id, _: send_for_review(callback_id, document, extracted_fields),
+        "Await Human review",
+        WaitForCallbackConfig(timeout=Duration.from_days(7)),
+    )
+
+    review_result = ReviewResult(**json.loads(review_result_str))
+
+    if not review_result.approved:
+        return {
+            "status": "rejected",
+            "notes": review_result.notes,
+            "extractedFields": extracted_fields,
+        }
+
+    return {"status": "approved", "extractedFields": extracted_fields}

lambda-durable-ai-human-review/example.ts

@@ -0,0 +1,62 @@
+import {
+  type DurableContext,
+  withDurableExecution,
+} from "@aws/durable-execution-sdk-js";
+import {
+  BedrockRuntimeClient,
+  ConverseCommand,
+} from "@aws-sdk/client-bedrock-runtime";
+import { z } from "zod";
+
+const MODEL_ID = "us.amazon.nova-pro-v1:0";
+const bedrock = new BedrockRuntimeClient({});
+
+async function converse(modelId: string, prompt: string): Promise<string> {
+  const response = await bedrock.send(
+    new ConverseCommand({
+      modelId,
+      messages: [{ role: "user", content: [{ text: prompt }] }],
+    }),
+  );
+  return response.output?.message?.content?.[0].text ?? "";
+}
+
+const ReviewResultSchema = z.object({
+  approved: z.boolean(),
+  notes: z.string().optional(),
+});
+
+const sendForReview = (
+  callbackId: string,
+  document: string,
+  extractedFields: string
+) => {
+  console.log(`Review needed for document. Callback ID: ${callbackId}`);
+  console.log(`Extracted fields: ${extractedFields}`);
+};
+
+export const handler = withDurableExecution(
+  async (event: { document?: string }, context: DurableContext) => {
+    const document = event.document ?? "Sample invoice with amount $1,234.56";
+
+    const extractedFields = await context.step("extract fields", async () =>
+      converse(MODEL_ID, `Extract key fields from this document as JSON: "${document}"`)
+    );
+
+    const reviewResultStr = await context.waitForCallback<string>(
+      "Await Human review",
+      async (callbackId) => {
+        sendForReview(callbackId, document, extractedFields);
+      },
+      { timeout: { days: 7 } }
+    );
+
+    const reviewResult = ReviewResultSchema.parse(JSON.parse(reviewResultStr));
+
+    if (!reviewResult.approved) {
+      return { status: "rejected", notes: reviewResult.notes, extractedFields };
+    }
+
+    return { status: "approved", extractedFields };
+  }
+);