The State Management Challenge for LLM Applications
Traditional ML models are stateless: each prediction is independent, with no memory of previous inputs. LLM applications are fundamentally different—they maintain conversation state across multiple turns, accumulating context that shapes every subsequent response. A customer support chatbot must remember that the user mentioned order #12345 three messages ago. A coding assistant must track which files have been discussed and what changes were already made. A research assistant must maintain a coherent understanding of a multi-document investigation spanning dozens of messages.
This statefulness creates MLOps challenges that don't exist for traditional ML: conversation storage at scale, context window management, memory costs that grow with conversation length, consistency across distributed deployments, and the need to evaluate multi-turn conversation quality. This article provides a complete architecture for managing state in production LLM applications.
Phase 1: Conversation State Architecture
State Storage Design
# state/conversation_store.py — Scalable conversation state management
import json
import hashlib
import time
from datetime import datetime, timezone
from typing import List, Dict, Optional, Any
from dataclasses import dataclass, asdict, field
from enum import Enum
import boto3
from boto3.dynamodb.conditions import Key
class MessageRole(Enum):
SYSTEM = "system"
USER = "user"
ASSISTANT = "assistant"
TOOL = "tool"
@dataclass
class Message:
role: str
content: str
timestamp: str = ""
metadata: Dict[str, Any] = field(default_factory=dict)
token_count: int = 0
message_id: str = ""
def __post_init__(self):
if not self.timestamp:
self.timestamp = datetime.now(timezone.utc).isoformat()
if not self.message_id:
self.message_id = hashlib.sha256(
f"{self.timestamp}:{self.content[:100]}".encode()
).hexdigest()[:16]
@dataclass
class ConversationState:
conversation_id: str
user_id: str
messages: List[Message] = field(default_factory=list)
context: Dict[str, Any] = field(default_factory=dict)
created_at: str = ""
updated_at: str = ""
total_tokens: int = 0
metadata: Dict[str, Any] = field(default_factory=dict)
def __post_init__(self):
if not self.created_at:
self.created_at = datetime.now(timezone.utc).isoformat()
self.updated_at = datetime.now(timezone.utc).isoformat()
class ConversationStore:
"""
Scalable conversation state storage using DynamoDB + S3.
DynamoDB stores conversation metadata and recent messages (fast access).
S3 stores full conversation history (cost-effective long-term storage).
"""
def __init__(self, table_name: str, s3_bucket: str):
self.dynamodb = boto3.resource("dynamodb")
self.table = self.dynamodb.Table(table_name)
self.s3 = boto3.client("s3")
self.s3_bucket = s3_bucket
def save_message(self, conversation_id: str, user_id: str,
message: Message) -> ConversationState:
"""Save a new message to a conversation."""
# Update DynamoDB record
response = self.table.update_item(
Key={
"user_id": user_id,
"conversation_id": conversation_id,
},
UpdateExpression="""
SET updated_at = :now,
total_tokens = total_tokens + :tokens,
message_count = if_not_exists(message_count, :zero) + :one,
#msgs = list_append(if_not_exists(#msgs, :empty), :new_msg)
""",
ExpressionAttributeNames={"#msgs": "recent_messages"},
ExpressionAttributeValues={
":now": datetime.now(timezone.utc).isoformat(),
":tokens": message.token_count,
":zero": 0,
":one": 1,
":new_msg": [asdict(message)],
":empty": [],
},
ReturnValues="ALL_NEW",
)
# If conversation has many messages, archive old ones to S3
item = response["Attributes"]
message_count = int(item.get("message_count", 0))
if message_count > 100: # Threshold for S3 archival
self._archive_old_messages(conversation_id, user_id)
return self._build_state(item)
def get_conversation(self, conversation_id: str, user_id: str,
max_messages: int = None) -> ConversationState:
"""Retrieve conversation state with optional message limit."""
response = self.table.get_item(
Key={"user_id": user_id, "conversation_id": conversation_id}
)
if "Item" not in response:
return ConversationState(
conversation_id=conversation_id,
user_id=user_id,
)
item = response["Item"]
state = self._build_state(item)
# If more messages needed, load from S3
if max_messages and len(state.messages) < max_messages:
archived = self._load_archived_messages(conversation_id, user_id)
state.messages = archived + state.messages
if max_messages:
state.messages = state.messages[-max_messages:]
return state
def get_context_window(self, conversation_id: str, user_id: str,
max_tokens: int = 4096) -> List[Message]:
"""
Get messages that fit within the context window token budget.
Uses a sliding window from most recent to oldest.
"""
state = self.get_conversation(conversation_id, user_id)
selected = []
token_budget = max_tokens
# Work backwards from most recent
for msg in reversed(state.messages):
if msg.token_count <= token_budget:
selected.insert(0, msg)
token_budget -= msg.token_count
else:
break
return selected
def _archive_old_messages(self, conversation_id: str, user_id: str):
"""Move old messages from DynamoDB to S3 for cost efficiency."""
item = self.table.get_item(
Key={"user_id": user_id, "conversation_id": conversation_id}
).get("Item", {})
messages = item.get("recent_messages", [])
if len(messages) <= 50:
return
# Archive all but the last 50 messages
to_archive = messages[:-50]
to_keep = messages[-50:]
# Upload to S3
s3_key = f"conversations/{user_id}/{conversation_id}/archive_{int(time.time())}.json"
self.s3.put_object(
Bucket=self.s3_bucket,
Key=s3_key,
Body=json.dumps(to_archive),
)
# Update DynamoDB to keep only recent messages
self.table.update_item(
Key={"user_id": user_id, "conversation_id": conversation_id},
UpdateExpression="SET recent_messages = :msgs, archive_keys = list_append(if_not_exists(archive_keys, :empty), :key)",
ExpressionAttributeValues={
":msgs": to_keep,
":empty": [],
":key": [s3_key],
},
)
def _load_archived_messages(self, conversation_id: str, user_id: str) -> List[Message]:
"""Load archived messages from S3."""
item = self.table.get_item(
Key={"user_id": user_id, "conversation_id": conversation_id}
).get("Item", {})
archive_keys = item.get("archive_keys", [])
messages = []
for key in archive_keys:
response = self.s3.get_object(Bucket=self.s3_bucket, Key=key)
archived = json.loads(response["Body"].read())
messages.extend([Message(**m) for m in archived])
return messages
def _build_state(self, item: dict) -> ConversationState:
"""Build ConversationState from DynamoDB item."""
messages = [
Message(**m) for m in item.get("recent_messages", [])
]
return ConversationState(
conversation_id=item["conversation_id"],
user_id=item["user_id"],
messages=messages,
context=item.get("context", {}),
created_at=item.get("created_at", ""),
updated_at=item.get("updated_at", ""),
total_tokens=int(item.get("total_tokens", 0)),
metadata=item.get("metadata", {}),
)Phase 2: Context Window Management
Intelligent Context Selection
# state/context_manager.py — Manage context window for LLM calls
import tiktoken
from typing import List, Dict, Optional
from dataclasses import dataclass
@dataclass
class ContextWindow:
system_prompt: str
messages: List[Dict]
total_tokens: int
truncated: bool
summary: Optional[str] = None
class ContextWindowManager:
"""
Manage the context window for LLM API calls.
Handles token counting, conversation summarization, and sliding window.
"""
def __init__(self, model: str = "gpt-4o", max_context_tokens: int = 8192):
self.model = model
self.max_context_tokens = max_context_tokens
self.encoder = tiktoken.encoding_for_model(model)
# Reserve tokens for system prompt and response
self.reserved_tokens = 1500 # System prompt + response buffer
self.available_tokens = max_context_tokens - self.reserved_tokens
def build_context(self, system_prompt: str, messages: List[Message],
user_context: Dict = None) -> ContextWindow:
"""
Build a context window that fits within the token budget.
Strategy:
1. Always include system prompt
2. Always include the most recent messages
3. If older messages don't fit, summarize them
4. Include relevant user context (preferences, profile)
"""
system_tokens = self._count_tokens(system_prompt)
remaining_budget = self.available_tokens - system_tokens
# Build message list from most recent
api_messages = []
tokens_used = 0
included_count = 0
truncated = False
for msg in reversed(messages):
msg_tokens = self._count_tokens(msg.content) + 4 # 4 tokens for role formatting
if tokens_used + msg_tokens <= remaining_budget:
api_messages.insert(0, {
"role": msg.role,
"content": msg.content,
})
tokens_used += msg_tokens
included_count += 1
else:
truncated = True
break
# If truncated, add a summary of excluded messages
summary = None
if truncated and included_count < len(messages):
excluded_messages = messages[:len(messages) - included_count]
summary = self._summarize_messages(excluded_messages)
# Add summary as a system message
summary_tokens = self._count_tokens(summary)
if summary_tokens <= remaining_budget - tokens_used:
api_messages.insert(0, {
"role": "system",
"content": f"[Summary of earlier conversation]: {summary}",
})
tokens_used += summary_tokens
# Add user context if provided
if user_context:
context_str = self._format_user_context(user_context)
context_tokens = self._count_tokens(context_str)
if context_tokens <= remaining_budget - tokens_used:
api_messages.insert(0, {
"role": "system",
"content": f"[User context]: {context_str}",
})
tokens_used += context_tokens
return ContextWindow(
system_prompt=system_prompt,
messages=api_messages,
total_tokens=system_tokens + tokens_used + self.reserved_tokens,
truncated=truncated,
summary=summary,
)
def _count_tokens(self, text: str) -> int:
"""Count tokens in text using tiktoken."""
return len(self.encoder.encode(text))
def _summarize_messages(self, messages: List[Message]) -> str:
"""
Summarize older messages to preserve key context.
In production, this would call an LLM for summarization.
"""
# Simple extractive summary: first sentence of each message
summaries = []
for msg in messages[-20:]: # Summarize last 20 excluded messages
first_sentence = msg.content.split(".")[0][:200]
summaries.append(f"{msg.role}: {first_sentence}")
return " | ".join(summaries)
def _format_user_context(self, context: Dict) -> str:
"""Format user context for inclusion in system prompt."""
parts = []
if context.get("name"):
parts.append(f"User name: {context['name']}")
if context.get("plan"):
parts.append(f"Subscription plan: {context['plan']}")
if context.get("preferences"):
parts.append(f"Preferences: {json.dumps(context['preferences'])}")
if context.get("previous_topics"):
parts.append(f"Previously discussed: {', '.join(context['previous_topics'])}")
return "; ".join(parts)Phase 3: Distributed State Consistency
Redis-Based Session Cache
# state/distributed_cache.py — Distributed conversation cache with Redis
import redis
import json
from typing import Optional, List
from datetime import timedelta
class DistributedConversationCache:
"""
Redis-based cache for active conversations.
Provides sub-millisecond reads for the most recent conversation state.
"""
def __init__(self, redis_url: str, ttl_hours: int = 24):
self.redis = redis.from_url(redis_url, decode_responses=True)
self.ttl = timedelta(hours=ttl_hours)
def cache_conversation(self, conversation_id: str, user_id: str,
state: ConversationState):
"""Cache conversation state in Redis."""
key = f"conv:{user_id}:{conversation_id}"
# Store full state as JSON
self.redis.setex(
key,
self.ttl,
json.dumps({
"conversation_id": state.conversation_id,
"messages": [asdict(m) for m in state.messages[-50:]], # Last 50 messages
"context": state.context,
"total_tokens": state.total_tokens,
}),
)
# Also maintain a sorted set of active conversations per user
self.redis.zadd(
f"user_conversations:{user_id}",
{conversation_id: time.time()},
)
self.redis.expire(f"user_conversations:{user_id}", self.ttl)
def get_cached_conversation(self, conversation_id: str,
user_id: str) -> Optional[ConversationState]:
"""Get conversation from cache (fast path)."""
key = f"conv:{user_id}:{conversation_id}"
data = self.redis.get(key)
if data:
parsed = json.loads(data)
return ConversationState(
conversation_id=parsed["conversation_id"],
user_id=user_id,
messages=[Message(**m) for m in parsed["messages"]],
context=parsed["context"],
total_tokens=parsed["total_tokens"],
)
return None # Cache miss — fall back to DynamoDB
def get_active_conversations(self, user_id: str) -> List[str]:
"""Get list of active conversation IDs for a user."""
return self.redis.zrevrange(
f"user_conversations:{user_id}",
0, -1,
)
def invalidate(self, conversation_id: str, user_id: str):
"""Invalidate cached conversation."""
self.redis.delete(f"conv:{user_id}:{conversation_id}")
# Usage in the LLM serving layer
class StatefulLLMService:
"""LLM service with conversation state management."""
def __init__(self, conversation_store: ConversationStore,
cache: DistributedConversationCache,
context_manager: ContextWindowManager):
self.store = conversation_store
self.cache = cache
self.context_mgr = context_manager
async def chat(self, user_id: str, conversation_id: str,
user_message: str, user_context: Dict = None) -> str:
"""Process a chat message with full state management."""
# 1. Load conversation state (cache-first)
state = self.cache.get_cached_conversation(conversation_id, user_id)
if state is None:
state = self.store.get_conversation(conversation_id, user_id)
# 2. Add user message to state
user_msg = Message(role="user", content=user_message)
user_msg.token_count = self.context_mgr._count_tokens(user_message)
state.messages.append(user_msg)
state.total_tokens += user_msg.token_count
# 3. Build context window
context_window = self.context_mgr.build_context(
system_prompt=self._get_system_prompt(user_context),
messages=state.messages,
user_context=user_context,
)
# 4. Call LLM
response = await self._call_llm(
context_window.system_prompt,
context_window.messages,
)
# 5. Add assistant response to state
assistant_msg = Message(role="assistant", content=response)
assistant_msg.token_count = self.context_mgr._count_tokens(response)
state.messages.append(assistant_msg)
state.total_tokens += assistant_msg.token_count
# 6. Persist state (async — don't block response)
self.store.save_message(conversation_id, user_id, user_msg)
self.store.save_message(conversation_id, user_id, assistant_msg)
# 7. Update cache
self.cache.cache_conversation(conversation_id, user_id, state)
return responsePhase 4: Evaluation of Multi-Turn Conversations
Conversation Quality Metrics
# evaluation/conversation_eval.py — Evaluate multi-turn conversation quality
from typing import List, Dict
import json
class ConversationEvaluator:
"""Evaluate the quality of multi-turn LLM conversations."""
def evaluate_conversation(self, conversation: ConversationState,
rubric: Dict = None) -> Dict:
"""
Evaluate a conversation across multiple dimensions.
"""
metrics = {}
# 1. Coherence — does the conversation maintain logical flow?
metrics["coherence"] = self._evaluate_coherence(conversation.messages)
# 2. Context retention — does the model remember earlier context?
metrics["context_retention"] = self._evaluate_context_retention(
conversation.messages
)
# 3. Consistency — are responses consistent with each other?
metrics["consistency"] = self._evaluate_consistency(conversation.messages)
# 4. Completeness — does the model address all user questions?
metrics["completeness"] = self._evaluate_completeness(conversation.messages)
# 5. Efficiency — does the model respond concisely?
metrics["efficiency"] = self._evaluate_efficiency(conversation.messages)
metrics["overall"] = sum(metrics.values()) / len(metrics)
return metrics
def _evaluate_context_retention(self, messages: List[Message]) -> float:
"""
Test if the model retains context from earlier in the conversation.
Uses reference questions that require earlier context to answer.
"""
# This would use an LLM-as-judge approach in production
# Example: ask "What was the order number I mentioned earlier?"
# and check if the response correctly references earlier context
pass
def _evaluate_coherence(self, messages: List[Message]) -> float:
"""Evaluate logical flow between consecutive messages."""
passConclusion
Stateful LLM applications require infrastructure that traditional ML serving doesn't need: conversation storage, context window management, distributed caching, and multi-turn evaluation. The architecture presented here uses DynamoDB for durable conversation storage, S3 for long-term archival, and Redis for low-latency access to active conversations. The context window manager intelligently selects which messages to include within the token budget, summarizing older messages when necessary. Distributed caching ensures sub-100ms response times even for conversations with hundreds of messages. Teams that invest in proper state management deliver more coherent, contextually aware LLM experiences while controlling costs through intelligent context window optimization.
