fix(context): scale replay cap with context window
This commit is contained in:
+15
-3
@@ -57,7 +57,11 @@ from nanobot.session.goal_state import (
|
||||
sustained_goal_active,
|
||||
)
|
||||
from nanobot.session.keys import UNIFIED_SESSION_KEY, session_key_for_channel
|
||||
from nanobot.session.manager import DEFAULT_REPLAY_MAX_MESSAGES, Session, SessionManager
|
||||
from nanobot.session.manager import (
|
||||
Session,
|
||||
SessionManager,
|
||||
replay_max_messages_for_context,
|
||||
)
|
||||
from nanobot.utils.document import extract_documents, reference_non_image_attachments
|
||||
from nanobot.utils.helpers import image_placeholder_text
|
||||
from nanobot.utils.helpers import truncate_text as truncate_text_fn
|
||||
@@ -201,7 +205,7 @@ class AgentLoop:
|
||||
timezone: str | None = None,
|
||||
session_ttl_minutes: int = 0,
|
||||
consolidation_ratio: float = 0.5,
|
||||
max_messages: int = DEFAULT_REPLAY_MAX_MESSAGES,
|
||||
max_messages: int = 0,
|
||||
hooks: list[AgentHook] | None = None,
|
||||
unified_session: bool = False,
|
||||
disabled_skills: list[str] | None = None,
|
||||
@@ -292,8 +296,11 @@ class AgentLoop:
|
||||
llm_wall_timeout_for_session=lambda sk: runner_wall_llm_timeout_s(self.sessions, sk),
|
||||
)
|
||||
self._unified_session = unified_session
|
||||
self._explicit_max_messages = max_messages > 0
|
||||
self._max_messages = (
|
||||
max_messages if max_messages > 0 else DEFAULT_REPLAY_MAX_MESSAGES
|
||||
max_messages
|
||||
if self._explicit_max_messages
|
||||
else replay_max_messages_for_context(self.context_window_tokens)
|
||||
)
|
||||
self._running = False
|
||||
self._mcp_servers = mcp_servers or {}
|
||||
@@ -422,6 +429,7 @@ class AgentLoop:
|
||||
self.runner.provider = provider
|
||||
self.subagents.set_provider(provider, model)
|
||||
self.consolidator.set_provider(provider, model, context_window_tokens)
|
||||
self._sync_replay_max_messages()
|
||||
self._provider_signature = snapshot.signature
|
||||
if publish_update and self._runtime_model_publisher is not None:
|
||||
self._runtime_model_publisher(
|
||||
@@ -435,6 +443,10 @@ class AgentLoop:
|
||||
)
|
||||
logger.info("Runtime model switched for next turn: {} -> {}", old_model, model)
|
||||
|
||||
def _sync_replay_max_messages(self) -> None:
|
||||
if not self._explicit_max_messages:
|
||||
self._max_messages = replay_max_messages_for_context(self.context_window_tokens)
|
||||
|
||||
def _refresh_provider_snapshot(self) -> None:
|
||||
if self._provider_snapshot_loader is None:
|
||||
return
|
||||
|
||||
@@ -438,6 +438,9 @@ class MyTool(Tool, ContextAware):
|
||||
setattr(self._runtime_state, key, value)
|
||||
if key == "model":
|
||||
self._runtime_state._active_preset = None
|
||||
sync_replay = getattr(self._runtime_state, "_sync_replay_max_messages", None)
|
||||
if key == "context_window_tokens" and callable(sync_replay):
|
||||
sync_replay()
|
||||
if key == "max_iterations" and hasattr(self._runtime_state, "_sync_subagent_runtime_limits"):
|
||||
self._runtime_state._sync_subagent_runtime_limits()
|
||||
self._audit("modify", f"{key}: {old!r} -> {value!r}")
|
||||
|
||||
@@ -27,7 +27,9 @@ from nanobot.utils.helpers import (
|
||||
from nanobot.utils.subagent_channel_display import scrub_subagent_announce_body
|
||||
|
||||
FILE_MAX_MESSAGES = 2000
|
||||
DEFAULT_REPLAY_MAX_MESSAGES = 500
|
||||
MIN_REPLAY_MAX_MESSAGES = 120
|
||||
REPLAY_TOKENS_PER_MESSAGE = 100
|
||||
DEFAULT_REPLAY_MAX_MESSAGES = FILE_MAX_MESSAGES
|
||||
_MESSAGE_TIME_PREFIX_RE = re.compile(r"^\[Message Time: [^\]]+\]\n?")
|
||||
_LOCAL_IMAGE_BREADCRUMB_RE = re.compile(r"^\[image: (?:/|~)[^\]]+\]\s*$")
|
||||
_TOOL_CALL_ECHO_RE = re.compile(r'^\s*(?:generate_image|message)\([^)]*\)\s*$')
|
||||
@@ -44,6 +46,15 @@ _FORK_VOLATILE_METADATA_KEYS = {
|
||||
}
|
||||
|
||||
|
||||
def replay_max_messages_for_context(context_window_tokens: int | None) -> int:
|
||||
if not context_window_tokens or context_window_tokens <= 0:
|
||||
return DEFAULT_REPLAY_MAX_MESSAGES
|
||||
return min(
|
||||
FILE_MAX_MESSAGES,
|
||||
max(MIN_REPLAY_MAX_MESSAGES, context_window_tokens // REPLAY_TOKENS_PER_MESSAGE),
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_assistant_replay_text(content: str) -> str:
|
||||
"""Remove internal replay artifacts that the model may have copied before.
|
||||
|
||||
|
||||
@@ -11,20 +11,31 @@ from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.providers.base import LLMResponse
|
||||
from nanobot.session.manager import DEFAULT_REPLAY_MAX_MESSAGES, Session
|
||||
from nanobot.providers.factory import ProviderSnapshot
|
||||
from nanobot.session.manager import (
|
||||
DEFAULT_REPLAY_MAX_MESSAGES,
|
||||
Session,
|
||||
replay_max_messages_for_context,
|
||||
)
|
||||
|
||||
DEFAULT_MAX_MESSAGES = DEFAULT_REPLAY_MAX_MESSAGES
|
||||
|
||||
|
||||
def _make_loop(tmp_path: Path, max_messages: int = DEFAULT_MAX_MESSAGES) -> AgentLoop:
|
||||
def _make_loop(
|
||||
tmp_path: Path,
|
||||
max_messages: int = 0,
|
||||
context_window_tokens: int = 200_000,
|
||||
) -> AgentLoop:
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
provider.generation.max_tokens = 4096
|
||||
return AgentLoop(
|
||||
bus=MessageBus(),
|
||||
provider=provider,
|
||||
workspace=tmp_path,
|
||||
model="test-model",
|
||||
max_messages=max_messages,
|
||||
context_window_tokens=context_window_tokens,
|
||||
)
|
||||
|
||||
|
||||
@@ -51,17 +62,26 @@ def _tool_round(call_id: str) -> list[dict]:
|
||||
|
||||
|
||||
class TestMaxMessagesInit:
|
||||
"""Verify AgentLoop stores the config value correctly."""
|
||||
"""Verify AgentLoop derives the internal replay cap correctly."""
|
||||
|
||||
def test_default_is_builtin_limit(self, tmp_path: Path) -> None:
|
||||
def test_context_formula(self) -> None:
|
||||
assert replay_max_messages_for_context(8_000) == 120
|
||||
assert replay_max_messages_for_context(32_768) == 327
|
||||
assert replay_max_messages_for_context(200_000) == DEFAULT_MAX_MESSAGES
|
||||
|
||||
def test_default_for_200k_context_reaches_file_cap(self, tmp_path: Path) -> None:
|
||||
loop = _make_loop(tmp_path)
|
||||
assert loop._max_messages == DEFAULT_MAX_MESSAGES
|
||||
|
||||
def test_default_scales_with_context_window(self, tmp_path: Path) -> None:
|
||||
loop = _make_loop(tmp_path, context_window_tokens=32_768)
|
||||
assert loop._max_messages == 327
|
||||
|
||||
def test_positive_value_stored(self, tmp_path: Path) -> None:
|
||||
loop = _make_loop(tmp_path, max_messages=25)
|
||||
assert loop._max_messages == 25
|
||||
|
||||
def test_zero_uses_builtin_limit(self, tmp_path: Path) -> None:
|
||||
def test_zero_uses_context_derived_limit(self, tmp_path: Path) -> None:
|
||||
loop = _make_loop(tmp_path, max_messages=0)
|
||||
assert loop._max_messages == DEFAULT_MAX_MESSAGES
|
||||
|
||||
@@ -70,6 +90,30 @@ class TestMaxMessagesInit:
|
||||
loop = _make_loop(tmp_path, max_messages=-5)
|
||||
assert loop._max_messages == DEFAULT_MAX_MESSAGES
|
||||
|
||||
def test_provider_refresh_resyncs_context_derived_limit(self, tmp_path: Path) -> None:
|
||||
old_provider = MagicMock()
|
||||
old_provider.get_default_model.return_value = "old-model"
|
||||
old_provider.generation.max_tokens = 4096
|
||||
new_provider = MagicMock()
|
||||
new_provider.generation.max_tokens = 4096
|
||||
loop = AgentLoop(
|
||||
bus=MessageBus(),
|
||||
provider=old_provider,
|
||||
workspace=tmp_path,
|
||||
model="old-model",
|
||||
context_window_tokens=32_768,
|
||||
provider_snapshot_loader=lambda: ProviderSnapshot(
|
||||
provider=new_provider,
|
||||
model="new-model",
|
||||
context_window_tokens=200_000,
|
||||
signature=("new-model",),
|
||||
),
|
||||
)
|
||||
|
||||
assert loop._max_messages == 327
|
||||
loop._refresh_provider_snapshot()
|
||||
assert loop._max_messages == DEFAULT_MAX_MESSAGES
|
||||
|
||||
|
||||
class TestGetHistoryWithMaxMessages:
|
||||
"""Verify get_history respects max_messages parameter."""
|
||||
|
||||
@@ -236,10 +236,12 @@ class TestModifyRestricted:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_modify_context_window_valid(self):
|
||||
tool = _make_tool()
|
||||
loop = _make_mock_loop(_sync_replay_max_messages=MagicMock())
|
||||
tool = _make_tool(runtime_state=loop)
|
||||
result = await tool.execute(action="set", key="context_window_tokens", value=131072)
|
||||
assert "Set context_window_tokens" in result
|
||||
assert tool._runtime_state.context_window_tokens == 131072
|
||||
assert loop.context_window_tokens == 131072
|
||||
loop._sync_replay_max_messages.assert_called_once_with()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_modify_none_value_for_restricted_int(self):
|
||||
|
||||
Reference in New Issue
Block a user