1
// src/handler.ts — Lambda handler for serverless agent
2

3
import { BedrockRuntimeClient, ConverseCommand } from "@aws-sdk/client-bedrock-runtime";
4
import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
5
import { DynamoDBDocumentClient, GetCommand, PutCommand } from "@aws-sdk/lib-dynamodb";
6
import { createClient } from "redis";
7

8
const bedrock = new BedrockRuntimeClient({ region: process.env.AWS_REGION });
9
const dynamo = DynamoDBDocumentClient.from(new DynamoDBClient({}));
10
const redis = createClient({ url: process.env.REDIS_URL });
11

12
// Warm-start: connect Redis once
13
let redisConnected = false;
14

15
export const handler = async (event: any) => {
16
  // Parse the incoming request
17
  const body = typeof event.body === "string" ? JSON.parse(event.body) : event.body;
18
  const { prompt, sessionId, config } = body;
19

20
  if (!prompt) {
21
    return respond(400, { error: "Missing 'prompt' in request body" });
22
  }
23

24
  const sessionKey = sessionId || `session:${Date.now()}:${Math.random().toString(36).slice(2, 10)}`;
25

26
  try {
27
    // 1. Try semantic cache
28
    const cached = await checkCache(prompt);
29
    if (cached) {
30
      return respond(200, {
31
        response: cached,
32
        sessionId: sessionKey,
33
        cached: true,
34
      });
35
    }
36

37
    // 2. Load session from DynamoDB
38
    const session = await loadSession(sessionKey);
39

40
    // 3. Call Bedrock
41
    const llmResponse = await callBedrock(prompt, session.context);
42

43
    // 4. Tool execution (if LLM requests tools)
44
    const toolResults = await executeTools(llmResponse.toolCalls);
45

46
    // 5. Final response
47
    const finalResponse = toolResults.length > 0
48
      ? await callBedrockWithTools(prompt, llmResponse.text, toolResults)
49
      : { text: llmResponse.text };
50

51
    // 6. Save session
52
    await saveSession(sessionKey, {
53
      context: [...(session?.context || []), { role: "user", content: prompt }, { role: "assistant", content: finalResponse.text }],
54
      lastActive: Date.now(),
55
    });
56

57
    return respond(200, {
58
      response: finalResponse.text,
59
      sessionId: sessionKey,
60
      cached: false,
61
    });
62
  } catch (error: any) {
63
    console.error("Agent error:", error);
64
    return respond(500, {
65
      error: "Internal agent error",
66
      detail: process.env.DEBUG ? error.message : undefined,
67
    });
68
  }
69
};
70

71
// ──── Helpers ────
72

73
async function callBedrock(prompt: string, context?: any[]) {
74
  const messages = [
75
    ...(context || []).map((m: any) => ({
76
      role: m.role,
77
      content: [{ text: m.content }],
78
    })),
79
    { role: "user", content: [{ text: prompt }] },
80
  ];
81

82
  const command = new ConverseCommand({
83
    modelId: process.env.MODEL_ID || "anthropic.claude-3-5-sonnet-20241022-v2:0",
84
    messages,
85
    inferenceConfig: { maxTokens: 4096 },
86
    toolConfig: {
87
      tools: agentTools,
88
    },
89
  });
90

91
  const result = await bedrock.send(command);
92
  const response = result.output?.message;
93

94
  return {
95
    text: response?.content?.find((c: any) => c.text)?.text || "",
96
    toolCalls: response?.content?.filter((c: any) => c.toolUse) || [],
97
  };
98
}
99

100
async function callBedrockWithTools(prompt: string, llmText: string, toolResults: any[]) {
101
  const messages = [
102
    { role: "user", content: [{ text: prompt }] },
103
    { role: "assistant", content: [{ text: llmText }] },
104
    {
105
      role: "user",
106
      content: toolResults.map((r: any) => ({
107
        toolResult: {
108
          toolUseId: r.toolUseId,
109
          content: [{ text: JSON.stringify(r.result) }],
110
        },
111
      })),
112
    },
113
  ];
114

115
  const result = await bedrock.send(new ConverseCommand({
116
    modelId: process.env.MODEL_ID || "anthropic.claude-3-5-sonnet-20241022-v2:0",
117
    messages,
118
    inferenceConfig: { maxTokens: 4096 },
119
  }));
120

121
  return {
122
    text: result.output?.message?.content?.find((c: any) => c.text)?.text || "",
123
  };
124
}
125

126
const agentTools = [
127
  {
128
    toolSpec: {
129
      name: "github_list_issues",
130
      description: "List issues from a GitHub repository",
131
      inputSchema: {
132
        json: {
133
          type: "object",
134
          properties: {
135
            owner: { type: "string", description: "Repository owner" },
136
            repo: { type: "string", description: "Repository name" },
137
            state: { type: "string", enum: ["open", "closed", "all"], description: "Issue state" },
138
          },
139
          required: ["owner", "repo"],
140
        },
141
      },
142
    },
143
  },
144
  // Add more tools here
145
];
146

147
async function executeTools(toolCalls: any[]): Promise<any[]> {
148
  if (!toolCalls.length) return [];
149

150
  return Promise.all(toolCalls.map(async (call: any) => {
151
    const { name, input } = call.toolUse;
152

153
    switch (name) {
154
      case "github_list_issues":
155
        return {
156
          toolUseId: call.toolUse.toolUseId,
157
          result: await listGitHubIssues(input.owner, input.repo, input.state),
158
        };
159
      default:
160
        return {
161
          toolUseId: call.toolUse.toolUseId,
162
          result: { error: `Unknown tool: ${name}` },
163
        };
164
    }
165
  }));
166
}
167

168
async function listGitHubIssues(owner: string, repo: string, state?: string): Promise<any> {
169
  const token = process.env.GITHUB_TOKEN;
170
  const url = `https://api.github.com/repos/${owner}/${repo}/issues?state=${state || "open"}&per_page=5`;
171

172
  const response = await fetch(url, {
173
    headers: {
174
      Authorization: `Bearer ${token}`,
175
      Accept: "application/vnd.github+json",
176
      "User-Agent": "serverless-agent",
177
    },
178
  });
179

180
  if (!response.ok) {
181
    return { error: `GitHub API error: ${response.status}` };
182
  }
183

184
  const issues = await response.json();
185
  return issues.map((i: any) => ({
186
    number: i.number,
187
    title: i.title,
188
    state: i.state,
189
    url: i.html_url,
190
  }));
191
}
192

193
async function checkCache(prompt: string): Promise<string | null> {
194
  if (shouldBypassCache(prompt)) return null;
195

196
  if (!redisConnected) {
197
    await redis.connect();
198
    redisConnected = true;
199
  }
200

201
  // Simplified cache check — see Day 3 for full semantic cache
202
  const cached = await redis.get(`exact:${simpleHash(prompt)}`);
203
  return cached;
204
}
205

206
async function loadSession(sessionId: string): Promise<any> {
207
  const result = await dynamo.send(new GetCommand({
208
    TableName: process.env.SESSION_TABLE || "agent-sessions",
209
    Key: { pk: `session#${sessionId}`, sk: "meta" },
210
  }));
211
  return result.Item?.data || null;
212
}
213

214
async function saveSession(sessionId: string, data: any): Promise<void> {
215
  await dynamo.send(new PutCommand({
216
    TableName: process.env.SESSION_TABLE || "agent-sessions",
217
    Item: {
218
      pk: `session#${sessionId}`,
219
      sk: "meta",
220
      data,
221
      ttl: Math.floor(Date.now() / 1000) + 3600,
222
    },
223
  }));
224
}
225

226
function shouldBypassCache(prompt: string): boolean {
227
  const writeKeywords = ["send", "create", "delete", "update", "write", "schedule"];
228
  return writeKeywords.some(k => prompt.toLowerCase().startsWith(k));
229
}
230

231
function simpleHash(s: string): string {
232
  let hash = 0;
233
  for (let i = 0; i < s.length; i++) {
234
    const char = s.charCodeAt(i);
235
    hash = ((hash << 5) - hash) + char;
236
    hash |= 0;
237
  }
238
  return Math.abs(hash).toString(16);
239
}
240

241
function respond(status: number, body: any) {
242
  return {
243
    statusCode: status,
244
    headers: {
245
      "Content-Type": "application/json",
246
      "Access-Control-Allow-Origin": "*",
247
      "Access-Control-Allow-Methods": "POST, OPTIONS",
248
      "Access-Control-Allow-Headers": "Content-Type, Authorization",
249
    },
250
    body: JSON.stringify(body),
251
  };
252
}

1
AWSTemplateFormatVersion: '2010-09-09'
2
Transform: AWS::Serverless-2016-10-31
3

4
Parameters:
5
  ModelId:
6
    Type: String
7
    Default: anthropic.claude-3-5-sonnet-20241022-v2:0
8
  GitHubToken:
9
    Type: String
10
    NoEcho: true
11
  RedisUrl:
12
    Type: String
13
  SessionTable:
14
    Type: String
15
    Default: agent-sessions
16

17
Globals:
18
  Function:
19
    Runtime: nodejs20.x
20
    Timeout: 30          # Lambda max for Bedrock calls
21
    MemorySize: 512
22
    Environment:
23
      Variables:
24
        MODEL_ID: !Ref ModelId
25
        GITHUB_TOKEN: !Ref GitHubToken
26
        REDIS_URL: !Ref RedisUrl
27
        SESSION_TABLE: !Ref SessionTable
28
    Policies:
29
      - Statement:
30
          - Effect: Allow
31
            Action:
32
              - bedrock:Converse
33
              - bedrock:InvokeModel
34
            Resource: "*"
35
          - Effect: Allow
36
            Action:
37
              - dynamodb:GetItem
38
              - dynamodb:PutItem
39
            Resource: !Sub "arn:aws:dynamodb:${AWS::Region}:${AWS::AccountId}:table/${SessionTable}"
40
          - Effect: Allow
41
            Action:
42
              - ec2:CreateNetworkInterface
43
              - ec2:DescribeNetworkInterfaces
44
              - ec2:DeleteNetworkInterface
45
            Resource: "*"  # Lambda VPC access for Redis
46

47
Resources:
48
  # ── Lambda Function ──
49
  AgentFunction:
50
    Type: AWS::Serverless::Function
51
    Properties:
52
      CodeUri: ./dist
53
      Handler: handler.handler
54
      Events:
55
        ApiEvent:
56
          Type: Api
57
          Properties:
58
            Path: /agent
59
            Method: POST
60
            RestApiId: !Ref AgentApi
61
        WebSocketEvent:
62
          Type: WebSocket
63
          Properties:
64
            Route: $default
65
            RouteResponseEnabled: true
66
            WebSocketApiId: !Ref AgentWebSocket
67

68
  # ── API Gateway REST ──
69
  AgentApi:
70
    Type: AWS::Serverless::Api
71
    Properties:
72
      StageName: prod
73
      OpenApiVersion: '3.0.3'
74
      Auth:
75
        DefaultAuthorizer: AWS_IAM
76
      EndpointConfiguration:
77
        Type: REGIONAL
78

79
  # ── API Gateway WebSocket ──
80
  AgentWebSocket:
81
    Type: AWS::Serverless::SimpleTable
82
    Properties:
83
      PrimaryKey:
84
        Name: connectionId
85
        Type: String
86

87
FunctionUrl:
88
    Type: AWS::Serverless::Function
89
    Properties:
90
      FunctionName: !Ref AgentFunction
91
      AutoPublishAlias: live
92
      Qualifier: live
93

94
Outputs:
95
  ApiEndpoint:
96
    Description: "HTTP API endpoint"
97
    Value: !Sub "https://${AgentApi}.execute-api.${AWS::Region}.amazonaws.com/prod/agent"
98
  WebSocketEndpoint:
99
    Description: "WebSocket endpoint"
100
    Value: !Sub "wss://${AgentWebSocket}.execute-api.${AWS::Region}.amazonaws.com/prod"
101
  FunctionUrl:
102
    Value: !GetAtt AgentFunctionUrl.FunctionUrl

1
export const wsHandler = async (event: any) => {
2
  const connectionId = event.requestContext.connectionId;
3
  const domain = event.requestContext.domainName;
4
  const stage = event.requestContext.stage;
5
  const callbackUrl = `https://${domain}/${stage}`;
6

7
  if (event.requestContext.routeKey === "$connect") {
8
    // Store connection
9
    await dynamo.send(new PutCommand({
10
      TableName: "agent-connections",
11
      Item: { connectionId, connectedAt: Date.now(), ttl: Math.floor(Date.now() / 1000) + 7200 },
12
    }));
13
    return { statusCode: 200 };
14
  }
15

16
  if (event.requestContext.routeKey === "$disconnect") {
17
    await dynamo.send(new DeleteCommand({
18
      TableName: "agent-connections",
19
      Key: { connectionId },
20
    }));
21
    return { statusCode: 200 };
22
  }
23

24
  // Handle message
25
  const body = JSON.parse(event.body);
26
  const { prompt, sessionId } = body;
27

28
  // Process asynchronously — send results back via WebSocket
29
  processAgentQuery(prompt, sessionId, connectionId, callbackUrl).catch(console.error);
30

31
  return { statusCode: 202 };  // Accepted — stream in progress
32
};