Pythoncontext_manager.py
Copy
1
import tiktoken
2
from typing import TypedDict
3
 
4
class Message(TypedDict):
5
    role: str
6
    content: str
7
 
8
# Model context limits (check docs for latest)
9
MODEL_LIMITS = {
10
    "gpt-4o": 128_000,
11
    "gpt-4o-mini": 128_000,
12
    "gpt-3.5-turbo": 16_385,
13
    "claude-3-5-sonnet-20241022": 200_000,
14
}
15
 
16
def count_tokens(messages: list[Message], model: str = "gpt-4o") -> int:
17
    """Count tokens in messages using tiktoken."""
18
    try:
19
        enc = tiktoken.encoding_for_model(model)
20
    except KeyError:
21
        enc = tiktoken.get_encoding("cl100k_base")
22
    
23
    # Per-message overhead (role tokens etc.)
24
    tokens = 3  # Priming tokens
25
    for msg in messages:
26
        tokens += 4  # message overhead
27
        tokens += len(enc.encode(msg["role"]))
28
        tokens += len(enc.encode(msg["content"]))
29
    
30
    return tokens
31
 
32
def truncate_messages(
33
    messages: list[Message],
34
    model: str = "gpt-4o",
35
    max_response_tokens: int = 4096,
36
    system_prompt: str | None = None,
37
) -> list[Message]:
38
    """
39
    Truncate messages to fit within context window.
40
    Preserves: system prompt (always) + most recent messages.
41
    """
42
    limit = MODEL_LIMITS.get(model, 128_000)
43
    budget = limit - max_response_tokens - 100  # Safety margin
44
    
45
    result: list[Message] = []
46
    
47
    # Always include system prompt
48
    if system_prompt:
49
        result.append({"role": "system", "content": system_prompt})
50
    
51
    # Add messages from most recent, working backwards
52
    non_system = [m for m in messages if m["role"] != "system"]
53
    included = []
54
    
55
    for msg in reversed(non_system):
56
        test = result + included + [msg]
57
        if count_tokens(test, model) <= budget:
58
            included.insert(0, msg)
59
        else:
60
            print(f"⚠️  Truncated {len(non_system) - len(included)} oldest messages")
61
            break
62
    
63
    result.extend(included)
64
    
65
    used = count_tokens(result, model)
66
    print(f"Context: {used:,} / {limit:,} tokens ({used/limit:.1%})")
67
    return result
68
 
69
# Example
70
history: list[Message] = [
71
    {"role": "user", "content": "Tell me about transformers."},
72
    {"role": "assistant", "content": "Transformers are neural network architectures..."},
73
    {"role": "user", "content": "How does attention work?"},
74
    {"role": "assistant", "content": "Attention computes weighted sums..."},
75
    {"role": "user", "content": "What about positional encoding?"},
76
]
77
 
78
trimmed = truncate_messages(
79
    history,
80
    model="gpt-4o",
81
    max_response_tokens=2048,
82
    system_prompt="You are a helpful ML tutor.",
83
)
84
print(f"Messages retained: {len(trimmed)}")
85