AWS Deployment

Use AWS Lambda for compute, API Gateway for WebSocket, Step Functions for durable execution, and RDS/Aurora for persistence.

Architecture

Browser ──WebSocket──▶ API Gateway (WebSocket API)
                            │
                            ▼
                      Lambda (connect/disconnect handler)
                            │
POST /messages ──▶ API Gateway (HTTP) ──▶ Lambda
                      │                       │
                      │                       ├─ Step Functions (durable execution)
                      │                       ├─ RDS/Aurora PostgreSQL
                      │                       └─ S3 Storage
                      │
                      └──▶ Response: { runId }

Setup

import { AI } from "@aimform/ai";
import { Pool } from "pg";
 
const pool = new Pool({ connectionString: process.env.DATABASE_URL });
 
const dbAdapter = {
  provider: "postgres",
  async query<T>(sql: string, params?: unknown[]): Promise<T[]> {
    const result = await pool.query(sql, params);
    return result.rows as T[];
  },
  async execute<T>(strings: TemplateStringsArray, ...values: unknown[]): Promise<T[]> {
    const sql = strings.reduce((acc, s, i) => acc + s + (i < values.length ? `$${i + 1}` : ""), "");
    return dbAdapter.query<T>(sql, values);
  },
};
 
// API Gateway WebSocket transport
const transport = new APIGatewayTransportAdapter(
  new ApiGatewayManagementApi({ endpoint: process.env.WS_ENDPOINT }),
);
 
const ai = new AI({
  model: "deepseek-v4-flash",
  apiKeys: { deepseekApiKey: process.env.DEEPSEEK_API_KEY! },
  db: dbAdapter as any,
  transport,
  enableMemory: true,
  plugins: {
    files: new S3FilePlugin(s3Client, "aimform-files"),
    webSearch: new TavilySearchPlugin(process.env.TAVILY_API_KEY!),
  },
});

Lambda Handler

import { Handler } from "aws-lambda";
 
export const handler: Handler = async (event) => {
  const { runId, messages, toolContext } = JSON.parse(event.body);
 
  const result = await ai.chat({
    runId,
    conversationId: toolContext.conversationId,
    messages,
    toolContext: { ...toolContext, env: process.env },
  });
 
  return { statusCode: 200, body: JSON.stringify(result) };
};

Step Functions

For long-running tasks (>30s Lambda timeout), use Step Functions:

{
  "StartAt": "Generate",
  "States": {
    "Generate": {
      "Type": "Task",
      "Resource": "arn:aws:lambda:...:function:ai-generate",
      "Next": "CheckTools"
    },
    "CheckTools": {
      "Type": "Choice",
      "Choices": [
        { "Variable": "$.hasTools", "BooleanEquals": true, "Next": "ExecuteTools" }
      ],
      "Default": "Finalize"
    },
    "ExecuteTools": {
      "Type": "Task",
      "Resource": "arn:aws:lambda:...:function:ai-execute",
      "Next": "Generate"
    },
    "Finalize": {
      "Type": "Task",
      "Resource": "arn:aws:lambda:...:function:ai-finalize",
      "End": true
    }
  }
}

WebSocket Transport (API Gateway)

class APIGatewayTransportAdapter implements ChatTransportAdapter {
  readonly provider = "aws-api-gateway";
 
  async publishEvents(runId: string, events: AnyChatEvent[]) {
    const connections = await getConnectionsByRunId(runId);
    for (const conn of connections) {
      await apiGateway.postToConnection({
        ConnectionId: conn.connectionId,
        Data: JSON.stringify({ events }),
      }).promise().catch(() => {});
    }
  }
 
  getConnectionUrl(runId: string): string {
    return `${process.env.WS_ENDPOINT}?runId=${runId}`;
  }
}