mirror of
https://github.com/usestrix/strix.git
synced 2026-09-11 22:51:18 +00:00
fix(llm): include system prompt tokens in memory compressor budget
The memory compressor was not accounting for system prompt and agent identity tokens when calculating the conversation budget. This caused premature history truncation on long scans with large system prompts. Adds a reserved_tokens parameter to compress_history() that subtracts already-accounted tokens from the budget before applying limits.
This commit is contained in:
parent
c9d2477144
commit
4f852e662f
2 changed files with 15 additions and 3 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue