The Hidden Cost of Context Windows: Managing Tokens in Production

import anthropic
 
client = anthropic.Anthropic()
 
def count_tokens(messages: list[dict], system: str = "") -> int:
    response = client.beta.messages.count_tokens(
        model="claude-3-5-sonnet-20241022",
        system=system,
        messages=messages,
    )
    return response.input_tokens
 
# In your request handler
def build_request(conversation: Conversation) -> dict:
    messages = conversation.to_messages()
    token_count = count_tokens(messages, system=SYSTEM_PROMPT)
 
    # Log for cost tracking
    logger.info("llm_request", tokens=token_count, conversation_id=conversation.id)
 
    if token_count > TOKEN_BUDGET:
        messages = compress_conversation(messages, target_tokens=TOKEN_BUDGET)
 
    return {"messages": messages, "model": MODEL}

def sliding_window(
    messages: list[dict],
    max_tokens: int,
    system_prompt: str,
    min_messages: int = 4,  # Always keep at least the last few turns
) -> list[dict]:
    if count_tokens(messages, system_prompt) <= max_tokens:
        return messages
 
    # Drop oldest messages (but preserve at least min_messages)
    while (
        len(messages) > min_messages
        and count_tokens(messages, system_prompt) > max_tokens
    ):
        # Remove oldest user+assistant pair
        messages = messages[2:]
 
    return messages

SUMMARISE_PROMPT = """
The following is the beginning of a longer conversation. 
Summarise it in 2–3 sentences, capturing any important context, 
decisions made, or facts established. Be concise.
 
Conversation:
{conversation}
 
Summary:"""
 
async def compress_conversation(
    messages: list[dict],
    target_tokens: int,
    system_prompt: str,
) -> list[dict]:
    # Split: old messages to summarise + recent messages to keep verbatim
    midpoint = len(messages) // 2
    to_summarise = messages[:midpoint]
    to_keep = messages[midpoint:]
 
    # Generate a summary of the old messages
    summary_text = await llm.complete(
        SUMMARISE_PROMPT.format(
            conversation=format_messages(to_summarise)
        )
    )
 
    # Replace old messages with a single summary message
    summary_message = {
        "role": "user",
        "content": f"[Earlier conversation summary: {summary_text}]",
    }
 
    compressed = [summary_message] + to_keep
 
    # Recurse if still over budget
    if count_tokens(compressed, system_prompt) > target_tokens:
        return await compress_conversation(compressed, target_tokens, system_prompt)
 
    return compressed

response = client.messages.create(
    model="claude-3-5-sonnet-20241022",
    max_tokens=1024,
    system=[
        {
            "type": "text",
            "text": LARGE_STATIC_SYSTEM_PROMPT,  # 10k tokens of medical guidelines
            "cache_control": {"type": "ephemeral"},  # Cache this block
        }
    ],
    messages=conversation_messages,
)
 
# Check cache performance
cache_hits = response.usage.cache_read_input_tokens
cache_misses = response.usage.cache_creation_input_tokens

from dataclasses import dataclass
from enum import Enum
 
class Feature(Enum):
    CHAT = "chat"
    CLINICAL_NOTE = "clinical_note"
    DOCUMENT_ANALYSIS = "document_analysis"
 
@dataclass
class TokenBudget:
    feature: Feature
    input_limit: int
    output_limit: int
    daily_limit_per_user: int
 
BUDGETS = {
    Feature.CHAT: TokenBudget(
        feature=Feature.CHAT,
        input_limit=16_000,
        output_limit=2_000,
        daily_limit_per_user=100_000,
    ),
    Feature.CLINICAL_NOTE: TokenBudget(
        feature=Feature.CLINICAL_NOTE,
        input_limit=8_000,
        output_limit=1_500,
        daily_limit_per_user=50_000,
    ),
}
 
async def check_budget(user_id: str, feature: Feature, estimated_tokens: int) -> bool:
    budget = BUDGETS[feature]
    used_today = await redis.get(f"tokens:{user_id}:{feature.value}:{today()}")
    return (used_today or 0) + estimated_tokens <= budget.daily_limit_per_user