Design an AI chatbot app

Question: Design an AI chatbot app

Source: https://www.hack2hire.com/question-bank/questions/system-design/69cb3259ff3ebde766cedc1f

Companies: Microsoft AI, Microsoft, Uber, OpenAI

Difficulty: Medium

Problem statement

Build a chatbot app on top of a hosted LLM API: conversations, streamed answers, history, and per-user usage limits.

Requirements

We are building a consumer chatbot app (web and mobile) that sits on top of a hosted LLM API: users hold multi-turn conversations, see answers stream in token by token, come back to their history later, and are held to per-user usage limits. We do not train or host the model; the provider's API is a dependency we pay for per token.

Functional requirements

• Sign up / sign in; every chat request is tied to an authenticated user and plan (free, plus).
• Create a conversation, send a message, and receive the assistant reply as a stream.
• List past conversations, open one, rename or delete it.
• Stop a reply mid-stream; regenerate the last reply.
• If the connection drops mid-answer, the client reconnects and gets the rest without paying twice.
• Per-user limits: messages per window and tokens per day by plan, with a clear "limit reached, resets at" response.

Non-functional requirements

• Time to first token under 1 s p50 on top of the provider's own latency; our overhead under 100 ms.
• A finished reply is never lost: once the user saw it, it is in history.
• Streams last 10–60 s, so the edge must hold many long-lived connections cheaply.
• Spend is bounded: no user, bug or retry loop can run up unbounded provider cost.
• History reads are eventually consistent across devices within a few seconds.

Back-of-the-envelope estimation

• 10M DAU, 15 messages/user/day → 150M messages/day ≈ 1,700/s average, ~5,000/s peak.
• Each reply streams ~20 s → 5,000/s × 20 s ≈ 100k concurrent open streams at peak. At ~20k idle-ish SSE connections per gateway node, that is ~5–10 nodes plus headroom; the provider's concurrency quota is the real ceiling.
• Tokens: ~2k input (history + prompt) + ~500 output per turn → 150M × 2.5k ≈ 375B tokens/day. Input dominates, which is why context trimming and prompt caching matter for cost.
• Storage: user message ~0.5 KB, reply ~2 KB → ~2.5 KB × 150M ≈ 375 GB/day of message text, ~135 TB/year before compression. Conversations: ~3M new/day.

API design

Send and stream (one request, the reply streams back as Server-Sent Events):

POST /v1/conversations/{conv_id}/messages
Authorization: Bearer <access token>
Idempotency-Key: <client uuid>
{ "content": "Explain Raft", "parent_message_id": "m_41", "model": "default" }

200 text/event-stream
event: message.created   data: {"user_message_id":"m_42","assistant_message_id":"m_43"}
event: delta   id: m_43:1   data: {"text":"Raft is "}
event: delta   id: m_43:2   data: {"text":"a consensus "}
...
event: done    id: m_43:87  data: {"finish_reason":"stop","usage":{"input":1830,"output":412}}
event: error   data: {"code":"moderation_blocked" | "provider_unavailable" | "quota_exceeded"}

• POST /v1/conversations → {conv_id}; GET /v1/conversations?cursor= → page of {conv_id, title, updated_at}.
• GET /v1/conversations/{id}/messages?before=&limit=50 → messages, newest page first.
• POST /v1/messages/{assistant_message_id}/cancel → stops generation; partial text is kept with status: cancelled.
• Recovery: GET /v1/messages/{assistant_message_id}/stream with Last-Event-ID: m_43:40 replays chunks 41+ and then tails live ones; if the message is already complete it returns the rest and done.
• Auth: POST /v1/auth/login → short-lived access token (15 min) + refresh token (httpOnly cookie); POST /v1/auth/refresh. A stream that outlives its token is allowed to finish (auth is checked at open, see deep dives).
• 429 body: {"code":"rate_limited","limit":"messages_3h","resets_at":"..."}.

Data model

Postgres (sharded by user_id) for users, conversations and messages: small relational rows, per-user access patterns, transactional writes.

Messages for one user live on one shard, so opening a conversation is a single-shard range scan. Very old conversations can move to cheaper storage (object store, loaded on open).

Redis for hot, short-lived state:

• stream:{assistant_message_id} — Redis Stream of chunks (seq, text), TTL 10 min after done. Used for replay on reconnect and for cross-node fan-out.
• quota:{user_id}:msgs:{window} — sliding-window counter for message limits.
• quota:{user_id}:tokens:{day} — reserved + spent tokens today, TTL 48 h.
• active:{user_id} — set of in-flight message ids, capped at 2–3 concurrent generations per user.

Storage trade-off: Redis holds only what the live stream needs; Postgres holds the finished text. We don't write every token to Postgres (that would be ~50 row updates per reply); we write once at done (or on cancel/failure), and Redis covers the gap.

High-level design

flowchart LR
  C[Web and mobile client] -->|"POST message, SSE"| GW[Edge gateway]
  GW --> AUTH[Auth service]
  GW --> CS[Chat service]
  CS --> Q[Quota service]
  Q --> R[(Redis)]
  CS --> CTX[Context builder]
  CTX --> DB[(Postgres shards)]
  CS --> MOD[Moderation]
  CS --> LLM[Hosted LLM API]
  CS -->|"chunks"| R
  CS --> DB
  CS --> USE[Usage ledger]
  USE --> DB

1. Send. The client POSTs with an idempotency key. The gateway validates the JWT locally (no auth round trip) and routes to a chat service node; the response stays open as SSE.
2. Admit. Chat service checks quotas in one Redis Lua call: message window, concurrent generations, and a token reservation (estimated input + max_output_tokens). Denied → 429 before any provider spend.
3. Persist the ask. Insert the user message and a placeholder assistant message (status: streaming) in one transaction; emit message.created so the client knows the ids it can resume by.
4. Build context. Load recent messages for the conversation, trim to the token budget (deep dive), add the system prompt, and run input moderation in parallel with the provider call's connection setup.
5. Stream. Call the provider with stream: true. For each provider chunk: append to stream:{id} in Redis with a sequence number, run it through the output moderation buffer, and write it to the SSE response.
6. Finish. On the provider's final event: write the full text and token counts to Postgres, settle the quota reservation with real usage, append to the usage ledger, send done.
7. History. Sidebar and conversation reads go straight to Postgres via the (user_id, updated_at) and (conv_id, message_id) indexes; a title is generated asynchronously after the first exchange with a small cheap model.

Why no message queue on the hot path. A queue between the request and the LLM call adds latency to time-to-first-token and breaks the natural backpressure of a held connection: the user is waiting on this exact request, and if it fails they retry. The only async work is fire-and-forget (title generation, analytics, ledger aggregation), which can go on a queue.

Deep dives

Streaming delivery and mid-stream recovery. SSE over HTTP/1.1 or HTTP/2 is the default: one-way, text, works through proxies, and the browser's EventSource already sends Last-Event-ID on reconnect. WebSockets only pay off if we add bidirectional features (voice, live collaboration).

The key decision is decoupling generation from the client connection. The provider call is driven by the chat service, not by the socket: if the client disconnects, generation continues to completion (bounded by max_tokens) and chunks keep landing in stream:{id}. On reconnect the client calls /stream with its last sequence number; any chat node reads XRANGE from that seq, then XREAD BLOCK for new chunks. The reply is billed once and persisted once. A heartbeat comment every 15 s keeps proxies from closing idle streams. If the chat node itself dies mid-generation, the placeholder stays streaming; a sweeper marks messages stuck for over 2 min as failed with the partial text, and the client offers "regenerate" (a new provider call, charged again, but only by explicit user action).

Auth-guarded streaming and connection lifecycle. Validate the access token once when the stream opens, then let the stream finish even if the token expires during it; re-checking per chunk buys nothing and costs latency. Revocation (logout everywhere, ban) is handled by a small revoked-session set the chat node checks every ~30 s for open streams, closing them with error: unauthorized. Recovery endpoints re-authenticate and check that the message belongs to the caller. Per-user caps on concurrent streams (3) and per-node caps on open connections (20k) stop one user or one node from exhausting file descriptors. On deploy, nodes drain: stop accepting new streams, let current ones finish for up to 60 s, then close; clients resume on another node via the Redis stream.

Context window and token budgeting. Budget = model window − max_output_tokens − system prompt − safety margin. Fill it newest-first: the latest user message always, then prior turns until the budget is hit. Long conversations get a rolling summary: when dropped turns exceed ~4k tokens, a background job summarizes them with a cheap model and stores summary + summary_upto_message_id on the conversation, which is prepended next time. Token counting uses the provider's tokenizer locally (estimate, then reconcile with the reported usage). Keeping the system prompt and summary as a stable prefix lets provider-side prompt caching cut input cost and latency, since input tokens are ~80% of volume.

Rate limiting and cost control. Two layers: a request limit (e.g. free: 40 messages/3 h; plus: 400) and a token budget per day, both in Redis. The token check is a reservation: reserve estimated_input + max_output before the call, settle to actual at done, release the remainder on cancel. That stops a burst of parallel requests from all passing a "spent < limit" check. A global spend breaker per model watches cost/minute across all users; if it exceeds plan, free-tier traffic is shed first or routed to a cheaper model. The usage_ledger in Postgres is the durable record; Redis counters are rebuilt from it if lost.

Streaming moderation and degradation. Input is moderated before or in parallel with the call; if it is flagged, the provider stream is aborted before tokens are shown. Output is checked on a sliding buffer: hold back ~100–200 characters, classify each window, and release text only after it passes, so a violation can be cut off with error: moderation_blocked and the partial reply is replaced, not just truncated. This adds a few hundred ms of visible lag, which is why the buffer stays small. When the provider is slow or erroring, retry only before the first token (with a fallback model or region); after tokens have streamed, fail the message instead of silently restarting with different text.

Trade-offs

• SSE over WebSockets keeps infrastructure plain HTTP but needs a separate request for cancel and gives up a single bidirectional channel.
• Generation continuing after disconnect guarantees the user gets the answer they paid for, at the cost of spending tokens on replies nobody may read; an explicit cancel is the escape hatch.
• Writing to Postgres once at done cuts write load ~50×, but a chat-node crash loses in-flight text not yet in Redis; Redis streams narrow that window.
• No queue on the hot path minimizes time-to-first-token and keeps backpressure natural, but bursts hit the provider directly; admission control and the spend breaker are the shock absorbers.
• Token reservations are accurate under concurrency but temporarily lock budget that a user may not end up using.
• Output moderation buffering improves safety at the cost of a short visible delay and some false-positive cut-offs.
• Rolling summaries keep long chats within budget and cheap, but lose detail from early turns.

Interviewer expectations

Mid-level (L4)

• Clear API for send/stream and history, with a sensible conversations/messages schema indexed by user and conversation.
• Chooses SSE or WebSockets with a reason and explains how tokens reach the browser incrementally.
• Basic per-user rate limiting in Redis and storing the full reply when it finishes.

Senior (L5)

• Decouples generation from the client connection and designs resume with sequence numbers and a replay buffer.
• Handles context-window trimming, token-based quotas with reservations, and why there is no queue on the hot path.
• Covers auth for long-lived streams, cancel, idempotency, and capacity math for concurrent connections.

Staff (L6+)

• Treats provider tokens as the dominant cost and designs global spend controls, model fallback, and prompt-cache-friendly context layout.
• Designs streaming moderation with its latency trade-off, graceful degradation and deploy draining for long streams.
• Defines observability per stream: time to first token, tokens/s, disconnect and resume rates, cost per user, with SLOs and alerts tied to them.