diff --git a/nanobot/agent/memory.py b/nanobot/agent/memory.py index fb08e371..56588bba 100644 --- a/nanobot/agent/memory.py +++ b/nanobot/agent/memory.py @@ -13,7 +13,6 @@ from datetime import datetime from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, Iterator -import tiktoken from loguru import logger from nanobot.session.manager import Session @@ -25,6 +24,7 @@ from nanobot.utils.helpers import ( find_legal_message_start, strip_think, truncate_text, + truncate_text_to_tokens, ) from nanobot.utils.prompt_templates import render_template @@ -806,14 +806,7 @@ class Consolidator: budget = self._input_token_budget if budget <= 0: return truncate_text(text, _RAW_ARCHIVE_MAX_CHARS) - try: - enc = tiktoken.get_encoding("cl100k_base") - tokens = enc.encode(text) - if len(tokens) <= budget: - return text - return enc.decode(tokens[:budget]) + "\n... (truncated)" - except Exception: - return truncate_text(text, budget * 4) + return truncate_text_to_tokens(text, budget) async def archive( self, diff --git a/nanobot/utils/helpers.py b/nanobot/utils/helpers.py index f6cfe774..01279931 100644 --- a/nanobot/utils/helpers.py +++ b/nanobot/utils/helpers.py @@ -218,6 +218,7 @@ _TOOL_RESULT_PREVIEW_CHARS = 1200 _TOOL_RESULTS_DIR = ".nanobot/tool-results" _TOOL_RESULT_RETENTION_SECS = 7 * 24 * 60 * 60 _TOOL_RESULT_MAX_BUCKETS = 32 +_TRUNCATED_SUFFIX = "\n... (truncated)" def safe_filename(name: str) -> str: @@ -234,7 +235,7 @@ def truncate_text(text: str, max_chars: int) -> str: """Truncate text with a stable suffix.""" if max_chars <= 0 or len(text) <= max_chars: return text - return text[:max_chars] + "\n... (truncated)" + return text[:max_chars] + _TRUNCATED_SUFFIX def truncate_text_to_tokens(text: str, max_tokens: int) -> str: @@ -252,9 +253,21 @@ def truncate_text_to_tokens(text: str, max_tokens: int) -> str: tokens = enc.encode(text) if len(tokens) <= max_tokens: return text - return enc.decode(tokens[:max_tokens]) + "\n... (truncated)" + suffix_tokens = enc.encode(_TRUNCATED_SUFFIX) + body_budget = max_tokens - len(suffix_tokens) + if body_budget <= 0: + return enc.decode(tokens[:max_tokens]) + result = enc.decode(tokens[:body_budget]) + _TRUNCATED_SUFFIX + while len(enc.encode(result)) > max_tokens and body_budget > 0: + body_budget -= 1 + result = enc.decode(tokens[:body_budget]) + _TRUNCATED_SUFFIX + return result except Exception: - return truncate_text(text, max_tokens * 4) + max_chars = max_tokens * 4 + suffix_chars = len(_TRUNCATED_SUFFIX) + if max_chars <= suffix_chars: + return text[:max_chars] + return truncate_text(text, max_chars - suffix_chars) def find_legal_message_start(messages: list[dict[str, Any]]) -> int: diff --git a/tests/agent/test_consolidator.py b/tests/agent/test_consolidator.py index 33754eb7..dafb1d40 100644 --- a/tests/agent/test_consolidator.py +++ b/tests/agent/test_consolidator.py @@ -825,4 +825,4 @@ class TestArchiveTruncation: enc = tiktoken.get_encoding("cl100k_base") sent_content = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"] token_count = len(enc.encode(sent_content)) - assert token_count <= 9_900 + 10 # small margin for truncation suffix + assert token_count <= 9_900 diff --git a/tests/agent/test_context_prompt_cache.py b/tests/agent/test_context_prompt_cache.py index 13f160ae..baed100c 100644 --- a/tests/agent/test_context_prompt_cache.py +++ b/tests/agent/test_context_prompt_cache.py @@ -237,8 +237,7 @@ def test_recent_history_truncated_at_max_tokens(tmp_path) -> None: assert len(history_section) == 2 enc = tiktoken.get_encoding("cl100k_base") - # Small margin for the truncation suffix appended after the token slice. - assert len(enc.encode(history_section[1])) <= builder._MAX_HISTORY_TOKENS + 50 + assert len(enc.encode(history_section[1])) <= builder._MAX_HISTORY_TOKENS def test_no_recent_history_when_dream_has_processed_all(tmp_path) -> None: diff --git a/tests/utils/test_helpers.py b/tests/utils/test_helpers.py index fd82ae05..bd6493df 100644 --- a/tests/utils/test_helpers.py +++ b/tests/utils/test_helpers.py @@ -24,8 +24,7 @@ def test_truncate_text_to_tokens_truncates_over_budget(): result = truncate_text_to_tokens(text, 50) assert result.endswith("\n... (truncated)") - body = result[: -len("\n... (truncated)")] - assert len(enc.encode(body)) <= 50 + assert len(enc.encode(result)) <= 50 def test_truncate_text_to_tokens_non_positive_budget_returns_text():