mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-04 16:38:49 +00:00
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:
parent
21d9072190
commit
072921893f
@ -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,
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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():
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user