Claude Function Calling on Bedrock: Wiring a Stateless Lambda for Real‑Time AI Actions

// src/lambda.ts
import { APIGatewayProxyEvent, APIGatewayProxyResult } from 'aws-lambda';
import {
  BedrockRuntimeClient,
  InvokeModelWithResponseStreamCommand,
  InvokeModelCommand,
} from '@aws-sdk/client-bedrock-runtime';
import { DynamoDBDocumentClient, PutCommand } from '@aws-sdk/lib-dynamodb';

// Create a Bedrock client that talks to the Claude model in the same region
const bedrock = new BedrockRuntimeClient({ region: process.env.AWS_REGION });

// Create a DynamoDB client that works with plain JavaScript objects
const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({ region: process.env.AWS_REGION }));

// The name of the DynamoDB table that stores users
const USER_TABLE = process.env.USER_TABLE!;

/**
 * Lambda entry point – receives an HTTP request from API Gateway,
 * forwards it to Claude, handles any tool call, writes to DynamoDB,
 * and returns a JSON response.
 */
export const handler = async (
  event: APIGatewayProxyEvent
): Promise<APIGatewayProxyResult> => {
  // 1️⃣ Parse the incoming body – we expect `{ "name": "Alice" }`
  const body = event.body ? JSON.parse(event.body) : {};

  // 2️⃣ Build the chat request with our tool definition (see next section)
  const chatPayload = buildChatPayload(body.name);

  // 3️⃣ Choose streaming or not based on a query param
  const useStream = event.queryStringParameters?.stream === 'true';

  // 4️⃣ Call Bedrock
  const response = useStream
    ? await invokeClaudeStream(chatPayload)
    : await invokeClaudeOnce(chatPayload);

  // 5️⃣ If Claude asked us to run a tool, handle it
  if (response.toolCall) {
    await handleToolCall(response.toolCall);
    // After the side‑effect we can return a simple confirmation
    return {
      statusCode: 200,
      body: JSON.stringify({ message: 'User added', user: response.toolCall.arguments }),
    };
  }

  // 6️⃣ No tool call – just forward Claude’s answer
  return {
    statusCode: 200,
    body: JSON.stringify({ answer: response.content }),
  };
};

// src/tools.ts
export const addUserTool = {
  name: 'add_user',
  description: 'Add a new user record to the Users table.',
  // JSON Schema describing the arguments the model should supply
  parameters: {
    type: 'object',
    properties: {
      userId: { type: 'string', description: 'A UUID for the new user' },
      name: { type: 'string', description: 'Full name of the user' },
      email: { type: 'string', format: 'email', description: 'User email address' },
    },
    required: ['userId', 'name', 'email'],
    additionalProperties: false,
  },
} as const;

// src/chatPayload.ts
import { addUserTool } from './tools';

export function buildChatPayload(name: string) {
  return {
    // The model we want to run – Claude‑3‑Sonnet‑20240229 is a good default
    modelId: 'anthropic.claude-3-sonnet-20240229-v1:0',
    // Enable streaming if you prefer; leave false for simple JSON
    body: JSON.stringify({
      // System prompt – explains the assistant’s role
      system: 'You are a helpful assistant that can add users to a database.',
      // The user message that triggered the request
      messages: [{ role: 'user', content: `Add a user named ${name}` }],
      // Tell Claude about the tool we expose
      tools: [addUserTool],
    }),
    // Content‑type required by Bedrock
    contentType: 'application/json',
    accept: 'application/json',
  };
}

// src/invokeClaude.ts
import {
  BedrockRuntimeClient,
  InvokeModelWithResponseStreamCommand,
  InvokeModelCommand,
} from '@aws-sdk/client-bedrock-runtime';

/**
 * Non‑streaming call – returns a single JSON object.
 */
export async function invokeClaudeOnce(payload: any) {
  const command = new InvokeModelCommand({
    modelId: payload.modelId,
    body: payload.body,
    contentType: payload.contentType,
    accept: payload.accept,
  });
  const { body } = await bedrock.send(command);
  const raw = Buffer.from(body as Uint8Array).toString('utf8');
  const parsed = JSON.parse(raw);

  // Bedrock puts the model’s answer under `completion` and, if a tool was called,
  // under `toolResult`. We normalise both shapes here.
  return normalizeBedrockResponse(parsed);
}

/**
 * Streaming call – Bedrock sends a series of SSE events.
 */
export async function invokeClaudeStream(payload: any) {
  const command = new InvokeModelWithResponseStreamCommand({
    modelId: payload.modelId,
    body: payload.body,
    contentType: payload.contentType,
    accept: payload.accept,
  });
  const response = await bedrock.send(command);

  // The response is a readable stream of Uint8Array chunks.
  // We need to concatenate them, split on newlines, and parse each SSE.
  const chunks: Uint8Array[] = [];
  for await (const chunk of response.body!) {
    chunks.push(chunk);
  }
  const raw = Buffer.concat(chunks).toString('utf8');

  // SSE format: each line starts with "data: " followed by JSON.
  const events = raw
    .split('\n')
    .filter((line) => line.startsWith('data: '))
    .map((line) => JSON.parse(line.replace(/^data: /, '')));

  // The last event usually contains the final model output.
  const final = events[events.length - 1];
  return normalizeBedrockResponse(final);
}

/**
 * Helper – turns both streaming and non‑streaming shapes into a unified object.
 */
function normalizeBedrockResponse(raw: any) {
  // Non‑streaming: `completion` holds the text, `toolResult` may exist
  if (raw.completion) {
    return {
      content: raw.completion,
      toolCall: raw.toolResult?.toolCall ?? null,
    };
  }
  // Streaming: `output` holds the same structure
  if (raw.output) {
    return {
      content: raw.output?.content?.[0]?.text ?? '',
      toolCall: raw.output?.toolResult?.toolCall ?? null,
    };
  }
  // Fallback – no known shape
  return { content: '', toolCall: null };
}

// src/types.ts
export interface AddUserArgs {
  userId: string;
  name: string;
  email: string;
}

// A discriminated union that can grow with more tools later
export type ToolCall =
  | { name: 'add_user'; arguments: AddUserArgs }
  // | { name: 'other_tool'; arguments: OtherArgs };

// src/handleTool.ts
import { PutCommand } from '@aws-sdk/lib-dynamodb';
import { ddb, USER_TABLE } from './lambda'; // re‑use the clients created earlier
import { ToolCall } from './types';

/**
 * Executes the side‑effect requested by Claude.
 * Right now we only support `add_user`, but the pattern scales.
 */
export async function handleToolCall(tool: ToolCall) {
  if (tool.name === 'add_user') {
    // 1️⃣ Prepare the item for DynamoDB
    const item = {
      PK: `USER#${tool.arguments.userId}`, // Partition key pattern
      SK: 'METADATA',
      name: tool.arguments.name,
      email: tool.arguments.email,
      createdAt: new Date().toISOString(),
    };

    // 2️⃣ Write the item – PutCommand replaces any existing record with the same PK.
    //    We could use TransactWriteItems for atomic multi‑item writes if needed.
    await ddb.send(new PutCommand({ TableName: USER_TABLE, Item: item }));
  } else {
    // Future tools can be added here
    throw new Error(`Unsupported tool: ${tool.name}`);
  }
}

# 1️⃣ Export required env vars
export AWS_REGION=us-east-1
export USER_TABLE=dev-Users
export AWS_ACCESS_KEY_ID=...
export AWS_SECRET_ACCESS_KEY=...

# 2️⃣ Run the Lambda locally (SAM will spin up a tiny API Gateway)
sam local invoke AddUserFunction -e events/add-user.json

{
  "body": "{\"name\":\"Bob\"}",
  "queryStringParameters": { "stream": "false" }
}

{
  "statusCode": 200,
  "body": "{\"message\":\"User added\",\"user\":{\"userId\":\"c9b1…\",\"name\":\"Bob\",\"email\":\"bob@example.com\"}}"
}

// src/__tests__/normalize.test.ts
import { normalizeBedrockResponse } from '../invokeClaude';

test('handles streaming shape', () => {
  const streamingRaw = {
    output: {
      content: [{ text: 'Done' }],
      toolResult: { toolCall: { name: 'add_user', arguments: { userId: '1', name: 'Bob', email: 'bob@example.com' } } },
    },
  };
  const result = normalizeBedrockResponse(streamingRaw);
  expect(result.toolCall?.name).toBe('add_user');
});

test('handles non‑streaming shape', () => {
  const nonStreamingRaw = {
    completion: 'Done',
    toolResult: { toolCall: { name: 'add_user', arguments: { userId: '2', name: 'Alice', email: 'alice@example.com' } } },
  };
  const result = normalizeBedrockResponse(nonStreamingRaw);
  expect(result.toolCall?.arguments.email).toBe('alice@example.com');
});