mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-09-04 02:01:48 +03:00
Code Issues Packages Projects Releases Wiki Activity

refactor(memory): remove consolidation ratio (#5575)

Browse Source 
* 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:
chengyongru
2026-08-28 13:20:09 +08:00
committed by GitHub
parent 29025f5a8b
commit cace42af14
8 changed files with 71 additions and 315 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/agent/test_consolidation_ratio.py
-112
View File
@@ -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)
tests/agent/test_consolidator.py
+7 -5
View File
@@ -232,17 +232,19 @@ class TestConsolidatorSummarize:
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)
assert "SNIP" in prompt
assert "final 4 conversation messages" in prompt
for mark in ("[permanent]", "[durable]", "[ephemeral]", "[correction]", "[skip]"):
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 mark something [skip] merely because it might already exist" in prompt
class TestConsolidatorArchiveErrorHandling:
"""archive() must fall back when the LLM does not complete its overview.
@@ -420,7 +422,7 @@ class TestConsolidatorTokenBudget:
consolidator.estimate_session_prompt_tokens = MagicMock(
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)
mock_provider.estimate_prompt_tokens.return_value = (100, "test-counter")
mock_provider.chat_with_retry.return_value = LLMResponse(
@@ -493,7 +495,7 @@ class TestConsolidatorTokenBudget:
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
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)
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
tests/agent/test_loop_consolidation_tokens.py
+20 -112
View File
@@ -2,7 +2,6 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
import nanobot.agent.memory as memory_module
from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus
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
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.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": role, "content": f"{role[0]}{turn}"}
for turn in range(10)
for role in ("user", "assistant")
]
loop.sessions.save(session)
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _message: 500)
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
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.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": role, "content": f"{role[0]}{turn}"}
for turn in range(10)
for role in ("user", "assistant")
]
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(
session,
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"]
archived_chunk = session.messages[:archive_end]
assert [message["content"] for message in archived_chunk] == ["u1", "a1", "u2", "a2"]
assert session.last_archived == 4
@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"},
assert [message["content"] for message in archived_chunk] == [
"u0", "a0", "u1", "a1", "u2", "a2", "u3", "a3", "u4", "a4", "u5", "a5",
]
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 (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
assert session.last_archived == 12
@pytest.mark.asyncio
async def test_consolidation_continues_below_trigger_until_half_target(tmp_path, monkeypatch) -> 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:
async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path) -> None:
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]
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": role, "content": f"{role[0]}{turn}"}
for turn in range(5)
for role in ("user", "assistant")
]
loop.sessions.save(session)
call_count = [0]
def mock_estimate(_session, *, runtime):
call_count[0] += 1
if call_count[0] == 1:
return (500, "test")
return (80, "test")
return (500, "test")
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(
session,
@@ -235,7 +145,7 @@ async def test_preflight_consolidation_receives_pending_summary(tmp_path) -> Non
@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."""
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.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": role, "content": f"{role[0]}{turn}"}
for turn in range(10)
for role in ("user", "assistant")
]
loop.sessions.save(session)
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 500)
call_count = [0]
def mock_estimate(_session, *, runtime):
call_count[0] += 1