fix: reuse token truncation helper
Maintainer edit: make token truncation include the suffix within the budget and route the consolidator through the shared helper so recent-history and archive truncation keep the same semantics.
This commit is contained in:
@@ -13,7 +13,6 @@ from datetime import datetime
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING, Any, Callable, Iterator
|
from typing import TYPE_CHECKING, Any, Callable, Iterator
|
||||||
|
|
||||||
import tiktoken
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from nanobot.session.manager import Session
|
from nanobot.session.manager import Session
|
||||||
@@ -25,6 +24,7 @@ from nanobot.utils.helpers import (
|
|||||||
find_legal_message_start,
|
find_legal_message_start,
|
||||||
strip_think,
|
strip_think,
|
||||||
truncate_text,
|
truncate_text,
|
||||||
|
truncate_text_to_tokens,
|
||||||
)
|
)
|
||||||
from nanobot.utils.prompt_templates import render_template
|
from nanobot.utils.prompt_templates import render_template
|
||||||
|
|
||||||
@@ -806,14 +806,7 @@ class Consolidator:
|
|||||||
budget = self._input_token_budget
|
budget = self._input_token_budget
|
||||||
if budget <= 0:
|
if budget <= 0:
|
||||||
return truncate_text(text, _RAW_ARCHIVE_MAX_CHARS)
|
return truncate_text(text, _RAW_ARCHIVE_MAX_CHARS)
|
||||||
try:
|
return truncate_text_to_tokens(text, budget)
|
||||||
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)
|
|
||||||
|
|
||||||
async def archive(
|
async def archive(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -218,6 +218,7 @@ _TOOL_RESULT_PREVIEW_CHARS = 1200
|
|||||||
_TOOL_RESULTS_DIR = ".nanobot/tool-results"
|
_TOOL_RESULTS_DIR = ".nanobot/tool-results"
|
||||||
_TOOL_RESULT_RETENTION_SECS = 7 * 24 * 60 * 60
|
_TOOL_RESULT_RETENTION_SECS = 7 * 24 * 60 * 60
|
||||||
_TOOL_RESULT_MAX_BUCKETS = 32
|
_TOOL_RESULT_MAX_BUCKETS = 32
|
||||||
|
_TRUNCATED_SUFFIX = "\n... (truncated)"
|
||||||
|
|
||||||
|
|
||||||
def safe_filename(name: str) -> str:
|
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."""
|
"""Truncate text with a stable suffix."""
|
||||||
if max_chars <= 0 or len(text) <= max_chars:
|
if max_chars <= 0 or len(text) <= max_chars:
|
||||||
return text
|
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:
|
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)
|
tokens = enc.encode(text)
|
||||||
if len(tokens) <= max_tokens:
|
if len(tokens) <= max_tokens:
|
||||||
return text
|
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:
|
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:
|
def find_legal_message_start(messages: list[dict[str, Any]]) -> int:
|
||||||
|
|||||||
@@ -825,4 +825,4 @@ class TestArchiveTruncation:
|
|||||||
enc = tiktoken.get_encoding("cl100k_base")
|
enc = tiktoken.get_encoding("cl100k_base")
|
||||||
sent_content = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"]
|
sent_content = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"]
|
||||||
token_count = len(enc.encode(sent_content))
|
token_count = len(enc.encode(sent_content))
|
||||||
assert token_count <= 9_900 + 10 # small margin for truncation suffix
|
assert token_count <= 9_900
|
||||||
|
|||||||
@@ -237,8 +237,7 @@ def test_recent_history_truncated_at_max_tokens(tmp_path) -> None:
|
|||||||
assert len(history_section) == 2
|
assert len(history_section) == 2
|
||||||
|
|
||||||
enc = tiktoken.get_encoding("cl100k_base")
|
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
|
||||||
assert len(enc.encode(history_section[1])) <= builder._MAX_HISTORY_TOKENS + 50
|
|
||||||
|
|
||||||
|
|
||||||
def test_no_recent_history_when_dream_has_processed_all(tmp_path) -> None:
|
def test_no_recent_history_when_dream_has_processed_all(tmp_path) -> None:
|
||||||
|
|||||||
@@ -24,8 +24,7 @@ def test_truncate_text_to_tokens_truncates_over_budget():
|
|||||||
result = truncate_text_to_tokens(text, 50)
|
result = truncate_text_to_tokens(text, 50)
|
||||||
|
|
||||||
assert result.endswith("\n... (truncated)")
|
assert result.endswith("\n... (truncated)")
|
||||||
body = result[: -len("\n... (truncated)")]
|
assert len(enc.encode(result)) <= 50
|
||||||
assert len(enc.encode(body)) <= 50
|
|
||||||
|
|
||||||
|
|
||||||
def test_truncate_text_to_tokens_non_positive_budget_returns_text():
|
def test_truncate_text_to_tokens_non_positive_budget_returns_text():
|
||||||
|
|||||||
Reference in New Issue
Block a user