Skip to content

Commit

Permalink
Discard old history items based on available tokens
Browse files Browse the repository at this point in the history
  • Loading branch information
muellerberndt committed Jan 2, 2025
1 parent e799ee6 commit b52c301
Showing 1 changed file with 18 additions and 1 deletion.
19 changes: 18 additions & 1 deletion src/ai/chatbot.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,24 @@ def _add_to_history(self, role: str, content: str) -> None:
# Keep history within limits, but preserve system message
if len(self.history) > self.max_history + 1: # +1 for system message
# Remove oldest messages but keep system message
self.history = [self.history[0]] + self.history[-(self.max_history) :]
self.history = [self.history[0]] + self.history[-(self.max_history):]

# Check total token usage and trim if needed
_, _, available = self.get_context_limits()
current_tokens = sum(self.count_tokens(msg["content"]) for msg in self.history)

if current_tokens > available:
# Always keep system message and last 2 messages of conversation
preserved = [self.history[0]] + self.history[-2:]
older_messages = self.history[1:-2]

# Remove older messages until we're under the limit
while current_tokens > available and older_messages:
removed = older_messages.pop()
current_tokens -= self.count_tokens(removed["content"])

# Reconstruct history with remaining messages
self.history = preserved if not older_messages else [self.history[0]] + older_messages + self.history[-2:]

def count_tokens(self, text: str) -> int:
"""Estimate token count for a string"""
Expand Down

0 comments on commit b52c301

Please sign in to comment.