Build a Serverless Bedrock AI Chatbot Like a Pro

This guide walks through a fully serverless chat assistant built on API Gateway → Lambda → Amazon Bedrock, with Cognito for authentication and a knowledge base powered by Amazon Titan embeddings + S3 Vectors for Retrieval-Augmented Generation (RAG). It also covers the part most tutorials skip: the cost engineering, model routing and prompt caching, that keeps a Claude-backed assistant affordable in production.

Architecture Overview

The request path is deliberately simple: the Browser calls API Gateway (used as a regular REST API), which invokes the Lambda, which calls Amazon Bedrock. Cognito validates the caller's token at API Gateway, so no authentication code lives in the Lambda. A separate RAG layer, Titan Text Embeddings plus an S3 Vectors index, supplies relevant context to the Lambda before it calls the Large Language Model (LLM). Generated tokens stream back along the same path, from Bedrock through the Lambda to the Browser, so the answer appears word by word instead of all at once.

Enable Streaming Response

An LLM answer takes seconds to generate. If you buffer the full response before returning it, the user watches a spinner for 5 to 20 seconds. If you stream tokens as they are produced, the first word appears in roughly 500 ms and the experience feels instant. The entire backend is therefore designed around not buffering.

Two pieces make server-side streaming work on AWS:

• Lambda response streaming: the function writes to a response stream instead of returning a single payload, using awslambda.streamifyResponse.
• API Gateway in streaming mode: a REST API integration configured with response_transfer_mode = "STREAM" against the Lambda's streaming invoke ARN, so bytes pass straight through instead of being collected.

The Lambda Handler

The handler uses the Node runtime with the Vercel AI SDK and the Bedrock provider. The key is streamifyResponse and piping the model's web stream into the Lambda response stream:

import { Readable } from 'node:stream'
import { pipeline } from 'node:stream/promises'
import { streamText, convertToModelMessages } from 'ai'
import { createAmazonBedrock } from '@ai-sdk/amazon-bedrock'

const bedrock = createAmazonBedrock({ region: process.env.BEDROCK_REGION })

export const handler = awslambda.streamifyResponse(async (event, responseStream) => {
  const { messages } = JSON.parse(event.body ?? '{}')

  const result = streamText({
    model: bedrock(process.env.BEDROCK_MODEL_ID),
    system: process.env.SYSTEM_PROMPT,
    messages: await convertToModelMessages(messages),
    maxOutputTokens: 2000
  })

  // Wrap the response with a 200 and streaming headers, then pipe the model's
  // token stream straight through, with no buffering.
  const http = awslambda.HttpResponseStream.from(responseStream, {
    statusCode: 200,
    headers: {
      'Content-Type': 'text/event-stream',
      'Cache-Control': 'no-cache, no-transform',
      'X-Accel-Buffering': 'no'
    }
  })

  await pipeline(Readable.fromWeb(result.toUIMessageStreamResponse().body), http)
})

The API Gateway Streaming Integration

The important part is response_transfer_mode = "STREAM" and pointing the URI at response_streaming_invoke_arn rather than the regular invoke_arn:

resource "aws_api_gateway_method" "chat_post" {
  rest_api_id   = aws_api_gateway_rest_api.this.id
  resource_id   = aws_api_gateway_resource.chat.id
  http_method   = "POST"
  authorization = "COGNITO_USER_POOLS"
  authorizer_id = aws_api_gateway_authorizer.cognito.id
}

resource "aws_api_gateway_integration" "chat_post" {
  rest_api_id             = aws_api_gateway_rest_api.this.id
  resource_id             = aws_api_gateway_resource.chat.id
  http_method             = aws_api_gateway_method.chat_post.http_method
  integration_http_method = "POST"
  type                    = "AWS_PROXY"

  # Streaming: pass tokens through as they arrive.
  uri                    = aws_lambda_function.chat.response_streaming_invoke_arn
  response_transfer_mode = "STREAM"
  timeout_milliseconds   = 900000   # streams can run long, up to 15 minutes
}

Putting it all together, here is the full lifecycle of a streaming chat request, from the browser through the Cognito-guarded API Gateway to Bedrock, with tokens flowing back token by token:

Securing the API with Cognito

If the assistant sits behind a login, let API Gateway validate the token for you with a Cognito User Pools authorizer, which keeps authentication code out of the Lambda entirely.

resource "aws_api_gateway_authorizer" "cognito" {
  name            = "chat-cognito"
  type            = "COGNITO_USER_POOLS"
  rest_api_id     = aws_api_gateway_rest_api.this.id
  provider_arns   = [aws_cognito_user_pool.main.arn]
  identity_source = "method.request.header.Authorization"
}

The SPA sends the Cognito ID token in the Authorization header, and API Gateway rejects anything unsigned with a 401 before your Lambda ever runs. The verified claims (sub, email, cognito:groups) arrive on event.requestContext.authorizer.claims, which is handy for per-user logic or admin-gating.

Grounding the Model with RAG

A raw LLM will confidently invent facts about your domain. Retrieval-Augmented Generation fixes that: you embed your knowledge base, then at query time retrieve the most relevant chunks and inject them into the prompt.

The classic pain point is the vector store. OpenSearch Serverless carries a non-trivial idle cost, whereas S3 Vectors is pay-per-use (storage plus queries) with no idle baseline, which is ideal for a small-to-medium knowledge base on AWS.

Provisioning the Vector Index

resource "aws_s3vectors_vector_bucket" "kb" {
  vector_bucket_name = "myapp-kb-vectors"
}

resource "aws_s3vectors_index" "kb" {
  vector_bucket_name = aws_s3vectors_vector_bucket.kb.vector_bucket_name
  index_name         = "docs"
  data_type          = "float32"
  dimension          = 1024      # must match the embedding model output
  distance_metric    = "cosine"

  metadata_configuration {
    non_filterable_metadata_keys = ["text"]   # store the chunk text alongside
  }
}

Indexing and Retrieval

There are two flows to keep separate:

• Indexing (on content save): chunk the document, embed each chunk with Amazon Titan Text Embeddings V2, and upsert the vectors keyed by ${docId}#${i}. Do this in your content-write path or a background job, never on the chat hot path.
• Retrieval (on each question): embed the question, query the index for the top-k nearest chunks, and prepend them to the system prompt.

import { BedrockRuntimeClient, InvokeModelCommand } from '@aws-sdk/client-bedrock-runtime'
import { S3VectorsClient, QueryVectorsCommand } from '@aws-sdk/client-s3vectors'

const runtime = new BedrockRuntimeClient({})
const vectors = new S3VectorsClient({})

async function retrieve(question) {
  // 1. Embed the question with Titan V2.
  const emb = await runtime.send(new InvokeModelCommand({
    modelId: 'amazon.titan-embed-text-v2:0',
    contentType: 'application/json',
    body: JSON.stringify({ inputText: question, dimensions: 1024, normalize: true })
  }))
  const { embedding } = JSON.parse(new TextDecoder().decode(emb.body))

  // 2. Nearest-neighbour search over the index.
  const res = await vectors.send(new QueryVectorsCommand({
    vectorBucketName: 'myapp-kb-vectors',
    indexName: 'docs',
    topK: 5,
    queryVector: { float32: embedding },
    returnMetadata: true,
    returnDistance: true          // keep the distances, they are reused for routing
  }))

  const chunks = res.vectors ?? []
  return {
    context: chunks.map(v => v.metadata.text).join('\n\n'),
    topDistance: chunks.length ? Math.min(...chunks.map(v => v.distance)) : null,
    count: chunks.length
  }
}

Fold context into the system prompt before calling streamText. Titan V2 embeddings are inexpensive (around $0.02 per million tokens) and embedding happens on save, so the RAG layer barely registers on the bill.

Controlling Cost

The architecture above works beautifully, and then the Bedrock invoice arrives. Frontier Claude models are billed per token, and output tokens cost roughly 5x input. As an EU inference-profile reference:

At roughly 2,000 input and 400 output tokens per message, a Sonnet-only assistant costs about $0.013 per message: trivial at 500 messages a month, but around $330 a month at 25,000. Two levers bring that down without gutting quality: model routing (send easy prompts to a cheap model, hard ones to a strong model) and prompt caching (stop paying full price for the repeated parts of the prompt). They stack, so let's take them in turn.

Model Routing

Managed: Bedrock Intelligent Prompt Routing

Bedrock has a built-in router: you point your request at a prompt router ARN instead of a model id, and Bedrock predicts, per request, which of two models in a family gives the best quality, routes there, and bills only the chosen model. Because your handler just passes a model id, adopting it is a one-line change (set BEDROCK_MODEL_ID to the router ARN) plus IAM to invoke the router:

# Invoke a default prompt router; the candidate models are covered by the
# broad foundation-model and inference-profile permissions.
statement {
  actions   = ["bedrock:InvokeModel", "bedrock:InvokeModelWithResponseStream"]
  resources = [
    "arn:aws:bedrock:*:*:default-prompt-router/*",
    "arn:aws:bedrock:*::foundation-model/*",
    "arn:aws:bedrock:*:*:inference-profile/*"
  ]
}

Application-Level Routing

Since your Lambda picks the model id per request, you can route between any models, including the latest ones. The only question is how you decide.

Heuristic Routing

Route on cheap, deterministic signals you already have. There is no extra model call and no added latency. The signals that work well for a RAG assistant are:

• Question length: long prompts tend to be complex.
• Conversation depth: deep threads carry more context to reason over.
• Grounding strength: reuse the topDistance from the retrieval step. A strong knowledge-base match means the cheap model can likely answer; a weak or missing match means the model must reason on its own, so escalate.
• Complexity keywords: "compare", "analyze", "debug", "why", "trade-offs".

const CHEAP  = 'eu.anthropic.claude-haiku-4-5-...'
const STRONG = 'eu.anthropic.claude-sonnet-4-5-...'

function chooseModel({ question, messageCount, topDistance, matchCount }) {
  const escalate =
    question.length > 600 ||
    messageCount > 6 ||
    matchCount === 0 ||
    (topDistance != null && topDistance > 0.6) ||   // weak grounding
    /\b(compare|analy[sz]e|evaluate|trade-?offs?|debug|architect|explain why)\b/i.test(question)

  return escalate ? STRONG : CHEAP
}

This approach is free, instant, transparent, and easy to tune (thresholds as environment variables) and observe (log the decision and signals per request). The downside is that it is a proxy for difficulty rather than a measure of it, so it will occasionally misroute a short-but-hard question. Tune the thresholds against real traffic.

Classifier Routing

Instead of hand-written rules, ask a tiny model to judge difficulty, then route on its verdict:

import { generateText } from 'ai'

async function classifyThenRoute(question) {
  const { text } = await generateText({
    model: bedrock('eu.amazon.nova-micro-...'),   // pennies per million tokens
    prompt: `Classify this request as "simple" or "complex". One word only.\n\n${question}`,
    maxOutputTokens: 1
  })
  return text.toLowerCase().includes('complex') ? STRONG : CHEAP
}

A classifier adapts to phrasing far better than a regex and catches the short-but-hard prompts a heuristic misses. The trade-off is one extra (cheap, fast) model call on every request, and the classifier itself can be wrong. Use a truly tiny model such as Nova Micro or Haiku so the overhead stays negligible.

Prompt Caching

Every request re-sends a lot of identical tokens: the system prompt, the retrieved knowledge chunks, and the growing conversation history. Bedrock prompt caching lets you mark a stable prefix with a cache point. On a hit within the TTL you pay a deep discount on those cached input tokens (roughly $0.30/M versus $3.30/M for Sonnet, about a 90% cut), while the first write costs a small premium.

You insert a cache breakpoint after the static prefix:

const result = streamText({
  model: bedrock(modelId),
  messages: [
    {
      role: 'system',
      content: `${SYSTEM_PROMPT}\n\n${context}`,
      // Mark everything above as cacheable.
      providerOptions: { bedrock: { cachePoint: { type: 'default' } } }
    },
    ...conversation
  ]
})

Read the fine print before you count the savings:

• Only input is cached. Output tokens, your priciest dimension, are never discounted. Since input and output cost are roughly equal per message at Sonnet rates, caching can cut at most about half the per-message cost.
• Short TTL (around 5 minutes). It shines for interactive, multi-turn chats where turns arrive quickly, and does nothing for one-shot questions minutes apart.
• Minimum cacheable size (around 1–2k tokens, depending on the model). Tiny prompts will not cache at all.
• RAG context changes per question, so it is mostly the system prompt and conversation history that cache reliably, not the freshly retrieved chunks.
• Caches are per-model, which is another reason to pin a conversation to one model.
• Verify SDK support for your @ai-sdk/amazon-bedrock version, as the exact cachePoint ergonomics vary.

Estimated Costs

Now that model routing and prompt caching are in place, let's look at their impact on cost across a few scenarios. Keeping the ~2,000 input / ~400 output tokens per message profile from earlier (a 5:1 input-to-output ratio), the three tiers below map to increasing monthly volume:

• Light — 1M input / 0.2M output (roughly 500 messages): a personal assistant, an internal tool for a handful of users, or a proof of concept.
• Normal — 10M input / 2M output (roughly 5,000 messages): a small production app with steady daily traffic.
• Heavy — 100M input / 20M output (roughly 50,000 messages): a busy, customer-facing assistant.

Reading down each column, the savings hold regardless of volume. Model routing alone brings the bill to about 57% of a Sonnet-only setup, roughly a 43% cut, by sending the easy two-thirds of traffic to Haiku (a third of Sonnet's price) and reserving Sonnet for the hard prompts. Adding prompt caching pushes total savings to around 55–60% on multi-turn workloads, where the repeated system prompt and history land in the cache. Concretely, the Normal tier drops from ~$66 to ~$37 with routing, then to ~$25–30 once caching kicks in.

Building an AI chatbot on Bedrock is easy, but keeping its cost under control while staying sustainable is too often an afterthought. As the numbers show, it pays to invest in features like model routing, prompt caching, output caps, and context trimming. If you take away one thing, make it model routing: it is the single biggest and most portable lever for any AI chatbot. It attacks the per-token price directly, the savings hold at every traffic level, and it applies to any provider or model family, not just Bedrock.

Wrapping Up

We have built a robust, fully serverless AI chatbot that covers the vast majority of real-world use cases: Cognito keeps it secure, RAG grounds answers in your own data, model routing and prompt caching keep the bill in check, and response streaming makes the experience feel instant.

A few more levers worth adding to improve your AI chatbot:

• Cap maxOutputTokens and trim your RAG topK and history to bound tokens.
• Add a token counter: track Bedrock's per-response token counts (e.g. in DynamoDB) and refuse new requests once a budget threshold is hit, so cost is enforced in real time instead of discovered on the invoice.