mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 08:13:11 +03:00
refactor(memory): remove consolidation ratio (#5575)
* refactor(memory): remove consolidation ratio * docs(memory): document fixed consolidation policy * docs(memory): simplify consolidation overview * docs(memory): rely on soft wrapping
This commit is contained in:
+1
-3
@@ -29,9 +29,7 @@ Memory moves through nanobot in two stages.
|
|||||||
|
|
||||||
### Stage 1: Consolidator
|
### Stage 1: Consolidator
|
||||||
|
|
||||||
When a conversation grows large enough to pressure the context window, nanobot does not try to carry every old message forever.
|
When a conversation grows large, the `Consolidator` summarizes older turns and appends the result to `memory/history.jsonl`, while keeping recent conversation available. Each summary preserves useful long-term facts and a short handoff for active work.
|
||||||
|
|
||||||
Instead, the `Consolidator` summarizes the oldest safe slice of the conversation and appends that summary to `memory/history.jsonl`.
|
|
||||||
|
|
||||||
This file is:
|
This file is:
|
||||||
|
|
||||||
|
|||||||
@@ -273,7 +273,6 @@ class AgentLoop:
|
|||||||
channels_config: ChannelsConfig | None = None,
|
channels_config: ChannelsConfig | None = None,
|
||||||
timezone: str | None = None,
|
timezone: str | None = None,
|
||||||
session_ttl_minutes: int = 0,
|
session_ttl_minutes: int = 0,
|
||||||
consolidation_ratio: float = 0.5,
|
|
||||||
hooks: list[AgentHook] | None = None,
|
hooks: list[AgentHook] | None = None,
|
||||||
hook_factories: list[AgentTurnHookFactory] | None = None,
|
hook_factories: list[AgentTurnHookFactory] | None = None,
|
||||||
unified_session: bool = False,
|
unified_session: bool = False,
|
||||||
@@ -444,7 +443,6 @@ class AgentLoop:
|
|||||||
workspace_scopes=self.workspace_scopes,
|
workspace_scopes=self.workspace_scopes,
|
||||||
unified_session=unified_session,
|
unified_session=unified_session,
|
||||||
),
|
),
|
||||||
consolidation_ratio=consolidation_ratio,
|
|
||||||
unified_session=unified_session,
|
unified_session=unified_session,
|
||||||
)
|
)
|
||||||
self.auto_compact = AutoCompact(
|
self.auto_compact = AutoCompact(
|
||||||
@@ -517,7 +515,6 @@ class AgentLoop:
|
|||||||
disabled_skills=defaults.disabled_skills,
|
disabled_skills=defaults.disabled_skills,
|
||||||
session_ttl_minutes=defaults.session_ttl_minutes,
|
session_ttl_minutes=defaults.session_ttl_minutes,
|
||||||
idle_compact_check_interval_seconds=defaults.idle_compact_check_interval_seconds,
|
idle_compact_check_interval_seconds=defaults.idle_compact_check_interval_seconds,
|
||||||
consolidation_ratio=defaults.consolidation_ratio,
|
|
||||||
tools_config=config.tools,
|
tools_config=config.tools,
|
||||||
model_presets=preset_helpers.configured_model_presets(config),
|
model_presets=preset_helpers.configured_model_presets(config),
|
||||||
model_preset=defaults.model_preset,
|
model_preset=defaults.model_preset,
|
||||||
|
|||||||
+41
-73
@@ -32,7 +32,6 @@ from nanobot.utils.gitstore import GitStore
|
|||||||
from nanobot.utils.helpers import (
|
from nanobot.utils.helpers import (
|
||||||
content_with_media_breadcrumbs,
|
content_with_media_breadcrumbs,
|
||||||
ensure_dir,
|
ensure_dir,
|
||||||
estimate_message_tokens,
|
|
||||||
estimate_prompt_tokens_chain,
|
estimate_prompt_tokens_chain,
|
||||||
strip_think,
|
strip_think,
|
||||||
truncate_text,
|
truncate_text,
|
||||||
@@ -956,8 +955,6 @@ class MemoryArchiver:
|
|||||||
class Consolidator:
|
class Consolidator:
|
||||||
"""Legacy context-pressure coordinator backed by a MemoryArchiver."""
|
"""Legacy context-pressure coordinator backed by a MemoryArchiver."""
|
||||||
|
|
||||||
_MAX_CONSOLIDATION_ROUNDS = 5
|
|
||||||
|
|
||||||
_SAFETY_BUFFER = 1024 # extra headroom for tokenizer estimation drift
|
_SAFETY_BUFFER = 1024 # extra headroom for tokenizer estimation drift
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -967,12 +964,10 @@ class Consolidator:
|
|||||||
build_messages: Callable[..., list[dict[str, Any]]],
|
build_messages: Callable[..., list[dict[str, Any]]],
|
||||||
get_tool_definitions: Callable[[], list[dict[str, Any]]],
|
get_tool_definitions: Callable[[], list[dict[str, Any]]],
|
||||||
resolve_prompt_context: Callable[[Session], tuple[str | None, Path | None]] | None = None,
|
resolve_prompt_context: Callable[[Session], tuple[str | None, Path | None]] | None = None,
|
||||||
consolidation_ratio: float = 0.5,
|
|
||||||
unified_session: bool = False,
|
unified_session: bool = False,
|
||||||
):
|
):
|
||||||
self.store = store
|
self.store = store
|
||||||
self.sessions = sessions
|
self.sessions = sessions
|
||||||
self.consolidation_ratio = consolidation_ratio
|
|
||||||
self.unified_session = unified_session
|
self.unified_session = unified_session
|
||||||
self._build_messages = build_messages
|
self._build_messages = build_messages
|
||||||
self._get_tool_definitions = get_tool_definitions
|
self._get_tool_definitions = get_tool_definitions
|
||||||
@@ -995,24 +990,19 @@ class Consolidator:
|
|||||||
def pick_consolidation_boundary(
|
def pick_consolidation_boundary(
|
||||||
self,
|
self,
|
||||||
session: Session,
|
session: Session,
|
||||||
tokens_to_remove: int,
|
) -> int | None:
|
||||||
) -> tuple[int, int] | None:
|
"""Return the fixed user-led boundary before the recent replay tail."""
|
||||||
"""Pick a user-turn boundary that removes enough old prompt tokens."""
|
if not session.messages:
|
||||||
start = session.last_archived
|
|
||||||
if start >= len(session.messages) or tokens_to_remove <= 0:
|
|
||||||
return None
|
return None
|
||||||
|
boundary = max(0, len(session.messages) - MIN_COMPACTED_REPLAY_MESSAGES)
|
||||||
removed_tokens = 0
|
while boundary > 0 and session.messages[boundary].get("role") != "user":
|
||||||
last_boundary: tuple[int, int] | None = None
|
boundary -= 1
|
||||||
for idx in range(start, len(session.messages)):
|
if (
|
||||||
message = session.messages[idx]
|
boundary <= session.last_archived
|
||||||
if idx > start and message.get("role") == "user":
|
or session.messages[boundary].get("role") != "user"
|
||||||
last_boundary = (idx, removed_tokens)
|
):
|
||||||
if removed_tokens >= tokens_to_remove:
|
return None
|
||||||
return last_boundary
|
return boundary
|
||||||
removed_tokens += estimate_message_tokens(message)
|
|
||||||
|
|
||||||
return last_boundary
|
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _full_replay_history(
|
def _full_replay_history(
|
||||||
@@ -1106,7 +1096,7 @@ class Consolidator:
|
|||||||
*,
|
*,
|
||||||
runtime: LLMRuntime,
|
runtime: LLMRuntime,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Loop: archive old messages until prompt fits within safe budget.
|
"""Archive one fixed old prefix when the prompt exceeds the safe budget.
|
||||||
|
|
||||||
The budget reserves space for completion tokens and a safety buffer
|
The budget reserves space for completion tokens and a safety buffer
|
||||||
so the LLM request never exceeds the context window.
|
so the LLM request never exceeds the context window.
|
||||||
@@ -1124,7 +1114,6 @@ class Consolidator:
|
|||||||
return
|
return
|
||||||
|
|
||||||
budget = self._input_token_budget(runtime)
|
budget = self._input_token_budget(runtime)
|
||||||
target = int(budget * self.consolidation_ratio)
|
|
||||||
last_summary: str | None = None
|
last_summary: str | None = None
|
||||||
estimated, source = self.estimate_session_prompt_tokens(
|
estimated, source = self.estimate_session_prompt_tokens(
|
||||||
session,
|
session,
|
||||||
@@ -1146,58 +1135,37 @@ class Consolidator:
|
|||||||
self._persist_last_summary(session, last_summary)
|
self._persist_last_summary(session, last_summary)
|
||||||
return
|
return
|
||||||
|
|
||||||
for round_num in range(self._MAX_CONSOLIDATION_ROUNDS):
|
end_idx = self.pick_consolidation_boundary(session)
|
||||||
if estimated <= target:
|
if end_idx is None:
|
||||||
break
|
logger.debug(
|
||||||
|
"Token consolidation: no safe fixed boundary for {}",
|
||||||
boundary = self.pick_consolidation_boundary(session, max(1, estimated - target))
|
|
||||||
if boundary is None:
|
|
||||||
logger.debug(
|
|
||||||
"Token consolidation: no safe boundary for {} (round {})",
|
|
||||||
session.key,
|
|
||||||
round_num,
|
|
||||||
)
|
|
||||||
break
|
|
||||||
|
|
||||||
end_idx = boundary[0]
|
|
||||||
|
|
||||||
chunk = session.messages[session.last_archived:end_idx]
|
|
||||||
if not chunk:
|
|
||||||
break
|
|
||||||
|
|
||||||
logger.info(
|
|
||||||
"Token consolidation round {} for {}: {}/{} via {}, chunk={} msgs",
|
|
||||||
round_num,
|
|
||||||
session.key,
|
session.key,
|
||||||
estimated,
|
|
||||||
runtime.context_window_tokens,
|
|
||||||
source,
|
|
||||||
len(chunk),
|
|
||||||
)
|
)
|
||||||
summary = await self.archive_session(
|
return
|
||||||
session,
|
|
||||||
archive_end=end_idx,
|
|
||||||
runtime=runtime,
|
|
||||||
)
|
|
||||||
# Advance the cursor either way: on success the chunk was
|
|
||||||
# summarized; on failure archive_session() raw-archived it as
|
|
||||||
# a breadcrumb. Re-archiving the same chunk on the next call
|
|
||||||
# would just emit duplicate [RAW] entries.
|
|
||||||
if summary:
|
|
||||||
last_summary = summary
|
|
||||||
session.last_archived = end_idx
|
|
||||||
self.sessions.save(session)
|
|
||||||
if not summary:
|
|
||||||
# LLM is degraded — stop hammering it this call;
|
|
||||||
# the next invocation can retry a fresh chunk.
|
|
||||||
break
|
|
||||||
|
|
||||||
estimated, source = self.estimate_session_prompt_tokens(
|
chunk = session.messages[session.last_archived:end_idx]
|
||||||
session,
|
if not chunk:
|
||||||
runtime=runtime,
|
return
|
||||||
)
|
|
||||||
if estimated <= 0:
|
logger.info(
|
||||||
break
|
"Token consolidation for {}: {}/{} via {}, chunk={} msgs",
|
||||||
|
session.key,
|
||||||
|
estimated,
|
||||||
|
runtime.context_window_tokens,
|
||||||
|
source,
|
||||||
|
len(chunk),
|
||||||
|
)
|
||||||
|
summary = await self.archive_session(
|
||||||
|
session,
|
||||||
|
archive_end=end_idx,
|
||||||
|
runtime=runtime,
|
||||||
|
)
|
||||||
|
# Advance either way: archive_session raw-archives on degradation,
|
||||||
|
# and replaying the same chunk would duplicate Memory material.
|
||||||
|
if summary:
|
||||||
|
last_summary = summary
|
||||||
|
session.last_archived = end_idx
|
||||||
|
self.sessions.save(session)
|
||||||
|
|
||||||
# Persist the last summary to session metadata so it can be injected
|
# Persist the last summary to session metadata so it can be injected
|
||||||
# into the runtime context on the next prepare_session() call, aligning
|
# into the runtime context on the next prepare_session() call, aligning
|
||||||
|
|||||||
@@ -155,13 +155,6 @@ class AgentDefaults(Base):
|
|||||||
default=60,
|
default=60,
|
||||||
ge=0,
|
ge=0,
|
||||||
) # Minimum interval in seconds between scans for idle sessions
|
) # Minimum interval in seconds between scans for idle sessions
|
||||||
consolidation_ratio: float = Field(
|
|
||||||
default=0.5,
|
|
||||||
ge=0.1,
|
|
||||||
le=0.95,
|
|
||||||
validation_alias=AliasChoices("consolidationRatio"),
|
|
||||||
serialization_alias="consolidationRatio",
|
|
||||||
) # Consolidation target ratio (0.5 = 50% of budget retained after compression)
|
|
||||||
dream: DreamConfig = Field(default_factory=DreamConfig)
|
dream: DreamConfig = Field(default_factory=DreamConfig)
|
||||||
|
|
||||||
@model_validator(mode="before")
|
@model_validator(mode="before")
|
||||||
|
|||||||
@@ -6,6 +6,8 @@ Use [skip] unless a fact meets all SNIP criteria:
|
|||||||
- Important: prevents rework or captures preferences / rules
|
- Important: prevents rework or captures preferences / rules
|
||||||
- Persistent: still relevant after 2 weeks
|
- Persistent: still relevant after 2 weeks
|
||||||
|
|
||||||
|
Also preserve a compact working-state handoff even when it is not Persistent: the active objective, current status, completed steps, unresolved blockers, next action, and exact identifiers needed to continue without rework. Mark these facts [ephemeral].
|
||||||
|
|
||||||
Format each fact as:
|
Format each fact as:
|
||||||
- [mark] fact content
|
- [mark] fact content
|
||||||
|
|
||||||
|
|||||||
@@ -1,112 +0,0 @@
|
|||||||
"""Tests for configurable consolidation_ratio."""
|
|
||||||
|
|
||||||
from unittest.mock import AsyncMock, MagicMock
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
from pydantic import ValidationError
|
|
||||||
|
|
||||||
import nanobot.agent.memory as memory_module
|
|
||||||
from nanobot.agent.loop import AgentLoop
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
from nanobot.config.schema import AgentDefaults
|
|
||||||
from nanobot.providers.base import GenerationSettings, LLMResponse
|
|
||||||
|
|
||||||
|
|
||||||
def _make_loop(
|
|
||||||
tmp_path,
|
|
||||||
*,
|
|
||||||
estimated_tokens: int = 0,
|
|
||||||
context_window_tokens: int = 200,
|
|
||||||
consolidation_ratio: float = 0.5,
|
|
||||||
) -> AgentLoop:
|
|
||||||
provider = MagicMock()
|
|
||||||
provider.get_default_model.return_value = "test-model"
|
|
||||||
provider.generation = GenerationSettings(max_tokens=0)
|
|
||||||
provider.estimate_prompt_tokens.return_value = (estimated_tokens, "test-counter")
|
|
||||||
_response = LLMResponse(content="ok", tool_calls=[])
|
|
||||||
provider.chat_with_retry = AsyncMock(return_value=_response)
|
|
||||||
provider.chat_stream_with_retry = AsyncMock(return_value=_response)
|
|
||||||
|
|
||||||
loop = AgentLoop(
|
|
||||||
bus=MessageBus(),
|
|
||||||
provider=provider,
|
|
||||||
workspace=tmp_path,
|
|
||||||
model="test-model",
|
|
||||||
context_window_tokens=context_window_tokens,
|
|
||||||
consolidation_ratio=consolidation_ratio,
|
|
||||||
)
|
|
||||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
|
||||||
loop.consolidator._SAFETY_BUFFER = 0
|
|
||||||
return loop
|
|
||||||
|
|
||||||
|
|
||||||
def _session_with_turns(loop: AgentLoop, *, turns: int):
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
|
||||||
session.messages = []
|
|
||||||
for i in range(turns):
|
|
||||||
session.messages.append({"role": "user", "content": f"u{i}", "timestamp": f"2026-01-01T00:00:{i:02d}"})
|
|
||||||
session.messages.append({"role": "assistant", "content": f"a{i}", "timestamp": f"2026-01-01T00:01:{i:02d}"})
|
|
||||||
loop.sessions.save(session)
|
|
||||||
return session
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
("ratio", "context_window_tokens", "estimates", "expected_archives"),
|
|
||||||
[
|
|
||||||
(0.5, 200, [250, 90], 1),
|
|
||||||
(0.1, 1000, [1200, 800, 400, 50], 2),
|
|
||||||
(0.9, 200, [300, 175], 1),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
async def test_consolidation_ratio_controls_target(
|
|
||||||
tmp_path,
|
|
||||||
monkeypatch,
|
|
||||||
ratio: float,
|
|
||||||
context_window_tokens: int,
|
|
||||||
estimates: list[int],
|
|
||||||
expected_archives: int,
|
|
||||||
) -> None:
|
|
||||||
loop = _make_loop(
|
|
||||||
tmp_path,
|
|
||||||
context_window_tokens=context_window_tokens,
|
|
||||||
consolidation_ratio=ratio,
|
|
||||||
)
|
|
||||||
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
|
|
||||||
session = _session_with_turns(loop, turns=10)
|
|
||||||
|
|
||||||
remaining_estimates = list(estimates)
|
|
||||||
|
|
||||||
runtime = loop.llm_runtime()
|
|
||||||
|
|
||||||
def mock_estimate(_session, *, runtime):
|
|
||||||
return (remaining_estimates.pop(0), "test")
|
|
||||||
|
|
||||||
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
|
|
||||||
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100)
|
|
||||||
|
|
||||||
await loop.consolidator.maybe_consolidate_by_tokens(
|
|
||||||
session,
|
|
||||||
runtime=runtime,
|
|
||||||
)
|
|
||||||
|
|
||||||
assert loop.consolidator.archive_session.await_count == expected_archives
|
|
||||||
|
|
||||||
|
|
||||||
def test_ratio_propagated_from_config_schema() -> None:
|
|
||||||
defaults = AgentDefaults()
|
|
||||||
assert defaults.consolidation_ratio == 0.5
|
|
||||||
|
|
||||||
defaults = AgentDefaults.model_validate({"consolidationRatio": 0.3})
|
|
||||||
assert defaults.consolidation_ratio == 0.3
|
|
||||||
|
|
||||||
dumped = defaults.model_dump(by_alias=True)
|
|
||||||
assert dumped["consolidationRatio"] == 0.3
|
|
||||||
|
|
||||||
|
|
||||||
def test_ratio_validation_rejects_out_of_range() -> None:
|
|
||||||
with pytest.raises(ValidationError):
|
|
||||||
AgentDefaults(consolidation_ratio=0.05)
|
|
||||||
|
|
||||||
with pytest.raises(ValidationError):
|
|
||||||
AgentDefaults(consolidation_ratio=1.0)
|
|
||||||
@@ -232,17 +232,19 @@ class TestConsolidatorSummarize:
|
|||||||
|
|
||||||
|
|
||||||
class TestConsolidatorPromptContract:
|
class TestConsolidatorPromptContract:
|
||||||
def test_archive_prompt_outputs_attribute_tags_without_missing_context_claims(self):
|
def test_archive_prompt_preserves_working_state_with_memory_facts(self):
|
||||||
prompt = render_template("agent/consolidator_archive.md", strip=True, archive_count=4)
|
prompt = render_template("agent/consolidator_archive.md", strip=True, archive_count=4)
|
||||||
|
|
||||||
assert "SNIP" in prompt
|
assert "SNIP" in prompt
|
||||||
assert "final 4 conversation messages" in prompt
|
assert "final 4 conversation messages" in prompt
|
||||||
for mark in ("[permanent]", "[durable]", "[ephemeral]", "[correction]", "[skip]"):
|
for mark in ("[permanent]", "[durable]", "[ephemeral]", "[correction]", "[skip]"):
|
||||||
assert mark in prompt
|
assert mark in prompt
|
||||||
assert "check context below" not in prompt.lower()
|
assert "working-state handoff" in prompt
|
||||||
|
assert "exact identifiers needed to continue without rework" in prompt
|
||||||
assert "Do not output facts already present in the system prompt's Recent History" in prompt
|
assert "Do not output facts already present in the system prompt's Recent History" in prompt
|
||||||
assert "Do not mark something [skip] merely because it might already exist" in prompt
|
assert "Do not mark something [skip] merely because it might already exist" in prompt
|
||||||
|
|
||||||
|
|
||||||
class TestConsolidatorArchiveErrorHandling:
|
class TestConsolidatorArchiveErrorHandling:
|
||||||
"""archive() must fall back when the LLM does not complete its overview.
|
"""archive() must fall back when the LLM does not complete its overview.
|
||||||
|
|
||||||
@@ -420,7 +422,7 @@ class TestConsolidatorTokenBudget:
|
|||||||
consolidator.estimate_session_prompt_tokens = MagicMock(
|
consolidator.estimate_session_prompt_tokens = MagicMock(
|
||||||
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
|
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
|
||||||
)
|
)
|
||||||
consolidator.pick_consolidation_boundary = MagicMock(return_value=(50, 800))
|
consolidator.pick_consolidation_boundary = MagicMock(return_value=50)
|
||||||
consolidator.archiver._build_messages = MagicMock(side_effect=_build_test_messages)
|
consolidator.archiver._build_messages = MagicMock(side_effect=_build_test_messages)
|
||||||
mock_provider.estimate_prompt_tokens.return_value = (100, "test-counter")
|
mock_provider.estimate_prompt_tokens.return_value = (100, "test-counter")
|
||||||
mock_provider.chat_with_retry.return_value = LLMResponse(
|
mock_provider.chat_with_retry.return_value = LLMResponse(
|
||||||
@@ -493,7 +495,7 @@ class TestConsolidatorTokenBudget:
|
|||||||
|
|
||||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||||
|
|
||||||
# Exactly one fallback per call — not _MAX_CONSOLIDATION_ROUNDS.
|
# The fixed policy archives at most one prefix per call.
|
||||||
assert consolidator.archive_session.await_count == 1
|
assert consolidator.archive_session.await_count == 1
|
||||||
|
|
||||||
async def test_boundary_respected_when_no_intermediate_user_turn(
|
async def test_boundary_respected_when_no_intermediate_user_turn(
|
||||||
@@ -520,7 +522,7 @@ class TestConsolidatorTokenBudget:
|
|||||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||||
|
|
||||||
consolidator.archive_session.assert_awaited_once()
|
consolidator.archive_session.assert_awaited_once()
|
||||||
# pick_consolidation_boundary finds the only boundary at idx=61
|
# The fixed recent tail expands backward to the user at idx=61.
|
||||||
assert session.last_archived == 61
|
assert session.last_archived == 61
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -2,7 +2,6 @@ from unittest.mock import AsyncMock, MagicMock
|
|||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
import nanobot.agent.memory as memory_module
|
|
||||||
from nanobot.agent.loop import AgentLoop
|
from nanobot.agent.loop import AgentLoop
|
||||||
from nanobot.bus.queue import MessageBus
|
from nanobot.bus.queue import MessageBus
|
||||||
from nanobot.providers.base import LLMResponse
|
from nanobot.providers.base import LLMResponse
|
||||||
@@ -41,17 +40,16 @@ async def test_prompt_below_threshold_does_not_consolidate(tmp_path) -> None:
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypatch) -> None:
|
async def test_prompt_above_threshold_triggers_consolidation(tmp_path) -> None:
|
||||||
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
|
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
|
||||||
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
|
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
session = loop.sessions.get_or_create("cli:test")
|
||||||
session.messages = [
|
session.messages = [
|
||||||
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
|
{"role": role, "content": f"{role[0]}{turn}"}
|
||||||
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
|
for turn in range(10)
|
||||||
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
|
for role in ("user", "assistant")
|
||||||
]
|
]
|
||||||
loop.sessions.save(session)
|
loop.sessions.save(session)
|
||||||
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _message: 500)
|
|
||||||
|
|
||||||
await loop.process_direct("hello", session_key="cli:test")
|
await loop.process_direct("hello", session_key="cli:test")
|
||||||
|
|
||||||
@@ -59,23 +57,18 @@ async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypat
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path, monkeypatch) -> None:
|
async def test_prompt_above_threshold_uses_fixed_recent_tail(tmp_path) -> None:
|
||||||
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
|
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
|
||||||
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
|
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
|
||||||
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
session = loop.sessions.get_or_create("cli:test")
|
||||||
session.messages = [
|
session.messages = [
|
||||||
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
|
{"role": role, "content": f"{role[0]}{turn}"}
|
||||||
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
|
for turn in range(10)
|
||||||
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
|
for role in ("user", "assistant")
|
||||||
{"role": "assistant", "content": "a2", "timestamp": "2026-01-01T00:00:03"},
|
|
||||||
{"role": "user", "content": "u3", "timestamp": "2026-01-01T00:00:04"},
|
|
||||||
]
|
]
|
||||||
loop.sessions.save(session)
|
loop.sessions.save(session)
|
||||||
|
|
||||||
token_map = {"u1": 120, "a1": 120, "u2": 120, "a2": 120, "u3": 120}
|
|
||||||
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda message: token_map[message["content"]])
|
|
||||||
|
|
||||||
await loop.consolidator.maybe_consolidate_by_tokens(
|
await loop.consolidator.maybe_consolidate_by_tokens(
|
||||||
session,
|
session,
|
||||||
runtime=loop.llm_runtime(),
|
runtime=loop.llm_runtime(),
|
||||||
@@ -83,112 +76,29 @@ async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path
|
|||||||
|
|
||||||
archive_end = loop.consolidator.archive_session.await_args.kwargs["archive_end"]
|
archive_end = loop.consolidator.archive_session.await_args.kwargs["archive_end"]
|
||||||
archived_chunk = session.messages[:archive_end]
|
archived_chunk = session.messages[:archive_end]
|
||||||
assert [message["content"] for message in archived_chunk] == ["u1", "a1", "u2", "a2"]
|
assert [message["content"] for message in archived_chunk] == [
|
||||||
assert session.last_archived == 4
|
"u0", "a0", "u1", "a1", "u2", "a2", "u3", "a3", "u4", "a4", "u5", "a5",
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> None:
|
|
||||||
"""Verify maybe_consolidate_by_tokens keeps looping until under threshold."""
|
|
||||||
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
|
|
||||||
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
|
|
||||||
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
|
||||||
session.messages = [
|
|
||||||
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
|
|
||||||
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
|
|
||||||
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
|
|
||||||
{"role": "assistant", "content": "a2", "timestamp": "2026-01-01T00:00:03"},
|
|
||||||
{"role": "user", "content": "u3", "timestamp": "2026-01-01T00:00:04"},
|
|
||||||
{"role": "assistant", "content": "a3", "timestamp": "2026-01-01T00:00:05"},
|
|
||||||
{"role": "user", "content": "u4", "timestamp": "2026-01-01T00:00:06"},
|
|
||||||
]
|
]
|
||||||
loop.sessions.save(session)
|
assert session.last_archived == 12
|
||||||
|
|
||||||
call_count = [0]
|
|
||||||
def mock_estimate(_session, *, runtime):
|
|
||||||
call_count[0] += 1
|
|
||||||
if call_count[0] == 1:
|
|
||||||
return (500, "test")
|
|
||||||
if call_count[0] == 2:
|
|
||||||
return (300, "test")
|
|
||||||
return (80, "test")
|
|
||||||
|
|
||||||
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
|
|
||||||
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100)
|
|
||||||
|
|
||||||
await loop.consolidator.maybe_consolidate_by_tokens(
|
|
||||||
session,
|
|
||||||
runtime=loop.llm_runtime(),
|
|
||||||
)
|
|
||||||
|
|
||||||
assert loop.consolidator.archive_session.await_count == 2
|
|
||||||
assert session.last_archived == 6
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_consolidation_continues_below_trigger_until_half_target(tmp_path, monkeypatch) -> None:
|
async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path) -> None:
|
||||||
"""Once triggered, consolidation should continue until it drops below half threshold."""
|
|
||||||
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
|
|
||||||
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
|
|
||||||
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
|
||||||
session.messages = [
|
|
||||||
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
|
|
||||||
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
|
|
||||||
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
|
|
||||||
{"role": "assistant", "content": "a2", "timestamp": "2026-01-01T00:00:03"},
|
|
||||||
{"role": "user", "content": "u3", "timestamp": "2026-01-01T00:00:04"},
|
|
||||||
{"role": "assistant", "content": "a3", "timestamp": "2026-01-01T00:00:05"},
|
|
||||||
{"role": "user", "content": "u4", "timestamp": "2026-01-01T00:00:06"},
|
|
||||||
]
|
|
||||||
loop.sessions.save(session)
|
|
||||||
|
|
||||||
call_count = [0]
|
|
||||||
|
|
||||||
def mock_estimate(_session, *, runtime):
|
|
||||||
call_count[0] += 1
|
|
||||||
if call_count[0] == 1:
|
|
||||||
return (500, "test")
|
|
||||||
if call_count[0] == 2:
|
|
||||||
return (150, "test")
|
|
||||||
return (80, "test")
|
|
||||||
|
|
||||||
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
|
|
||||||
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100)
|
|
||||||
|
|
||||||
await loop.consolidator.maybe_consolidate_by_tokens(
|
|
||||||
session,
|
|
||||||
runtime=loop.llm_runtime(),
|
|
||||||
)
|
|
||||||
|
|
||||||
assert loop.consolidator.archive_session.await_count == 2
|
|
||||||
assert session.last_archived == 6
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path, monkeypatch) -> None:
|
|
||||||
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
|
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
|
||||||
loop.consolidator.archive_session = AsyncMock(return_value="User discussed project status.") # type: ignore[method-assign]
|
loop.consolidator.archive_session = AsyncMock(return_value="User discussed project status.") # type: ignore[method-assign]
|
||||||
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
session = loop.sessions.get_or_create("cli:test")
|
||||||
session.messages = [
|
session.messages = [
|
||||||
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
|
{"role": role, "content": f"{role[0]}{turn}"}
|
||||||
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
|
for turn in range(5)
|
||||||
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
|
for role in ("user", "assistant")
|
||||||
]
|
]
|
||||||
loop.sessions.save(session)
|
loop.sessions.save(session)
|
||||||
|
|
||||||
call_count = [0]
|
|
||||||
|
|
||||||
def mock_estimate(_session, *, runtime):
|
def mock_estimate(_session, *, runtime):
|
||||||
call_count[0] += 1
|
return (500, "test")
|
||||||
if call_count[0] == 1:
|
|
||||||
return (500, "test")
|
|
||||||
return (80, "test")
|
|
||||||
|
|
||||||
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
|
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
|
||||||
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 150)
|
|
||||||
|
|
||||||
await loop.consolidator.maybe_consolidate_by_tokens(
|
await loop.consolidator.maybe_consolidate_by_tokens(
|
||||||
session,
|
session,
|
||||||
@@ -235,7 +145,7 @@ async def test_preflight_consolidation_receives_pending_summary(tmp_path) -> Non
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) -> None:
|
async def test_preflight_consolidation_before_llm_call(tmp_path) -> None:
|
||||||
"""Verify preflight consolidation runs before the LLM call in process_direct."""
|
"""Verify preflight consolidation runs before the LLM call in process_direct."""
|
||||||
order: list[str] = []
|
order: list[str] = []
|
||||||
|
|
||||||
@@ -258,13 +168,11 @@ async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) ->
|
|||||||
|
|
||||||
session = loop.sessions.get_or_create("cli:test")
|
session = loop.sessions.get_or_create("cli:test")
|
||||||
session.messages = [
|
session.messages = [
|
||||||
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
|
{"role": role, "content": f"{role[0]}{turn}"}
|
||||||
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
|
for turn in range(10)
|
||||||
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
|
for role in ("user", "assistant")
|
||||||
]
|
]
|
||||||
loop.sessions.save(session)
|
loop.sessions.save(session)
|
||||||
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 500)
|
|
||||||
|
|
||||||
call_count = [0]
|
call_count = [0]
|
||||||
def mock_estimate(_session, *, runtime):
|
def mock_estimate(_session, *, runtime):
|
||||||
call_count[0] += 1
|
call_count[0] += 1
|
||||||
|
|||||||
Reference in New Issue
Block a user