diff --git a/strix/llm/llm.py b/strix/llm/llm.py index fe1758f1..969710e1 100644 --- a/strix/llm/llm.py +++ b/strix/llm/llm.py @@ -10,7 +10,7 @@ from litellm.utils import supports_prompt_caching, supports_vision from strix.config import Config from strix.llm.config import LLMConfig -from strix.llm.memory_compressor import MemoryCompressor +from strix.llm.memory_compressor import MemoryCompressor, _get_message_tokens from strix.llm.utils import ( _truncate_to_first_function, fix_incomplete_tool_call, @@ -210,7 +210,12 @@ class LLM: } ) - compressed = list(self.memory_compressor.compress_history(conversation_history)) + reserved_tokens = sum( + _get_message_tokens(msg, self.config.litellm_model) for msg in messages + ) + compressed = list( + self.memory_compressor.compress_history(conversation_history, reserved_tokens) + ) conversation_history.clear() conversation_history.extend(compressed) messages.extend(compressed) diff --git a/strix/llm/memory_compressor.py b/strix/llm/memory_compressor.py index 8cad5107..728831b1 100644 --- a/strix/llm/memory_compressor.py +++ b/strix/llm/memory_compressor.py @@ -166,9 +166,16 @@ class MemoryCompressor: def compress_history( self, messages: list[dict[str, Any]], + reserved_tokens: int = 0, ) -> list[dict[str, Any]]: """Compress conversation history to stay within token limits. + Args: + messages: Conversation history messages to compress. + reserved_tokens: Tokens already reserved for system prompt and + other framing messages outside the conversation history. + Subtracted from the budget before checking limits. + Strategy: 1. Handle image limits first 2. Keep all system messages @@ -201,7 +208,7 @@ class MemoryCompressor: # Type assertion since we ensure model_name is not None in __init__ model_name: str = self.model_name # type: ignore[assignment] - total_tokens = sum( + total_tokens = reserved_tokens + sum( _get_message_tokens(msg, model_name) for msg in system_msgs + regular_msgs )