mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 00:03:01 +03:00
refactor(memory): unify structured consolidation flow
This commit is contained in:
@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Callable, Coroutine
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.session.manager import MIN_COMPACTED_REPLAY_MESSAGES, Session, SessionManager
|
||||
from nanobot.session.summary import SessionSummary
|
||||
from nanobot.session.summary import SessionSummary, session_summary_from_metadata
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from nanobot.agent.memory import Consolidator
|
||||
@@ -91,7 +91,7 @@ class AutoCompact:
|
||||
)
|
||||
if summary and summary != "(nothing)":
|
||||
session = self.sessions.get_or_create(key)
|
||||
stored = SessionSummary.from_metadata(
|
||||
stored = session_summary_from_metadata(
|
||||
session.metadata,
|
||||
fallback_last_active=session.updated_at,
|
||||
)
|
||||
@@ -117,7 +117,7 @@ class AutoCompact:
|
||||
# Cold path: summary persisted in session metadata (process restarted).
|
||||
# Persisted metadata may outlive schema changes; a malformed summary must
|
||||
# not abort turn preparation.
|
||||
return session, SessionSummary.from_metadata(
|
||||
return session, session_summary_from_metadata(
|
||||
session.metadata,
|
||||
fallback_last_active=session.updated_at,
|
||||
)
|
||||
|
||||
@@ -157,7 +157,11 @@ class ContextBuilder:
|
||||
parts.append("# Recent History\n\n" + history_text)
|
||||
|
||||
if session_summary:
|
||||
parts.append(f"[Archived Context Summary]\n\n{session_summary.for_prompt()}")
|
||||
parts.append(
|
||||
"[Archived Context Summary]\n\n"
|
||||
f"Previous conversation summary (last active {session_summary['last_active']}):\n"
|
||||
f"{session_summary['text']}"
|
||||
)
|
||||
|
||||
return "\n\n---\n\n".join(parts)
|
||||
|
||||
@@ -175,7 +179,7 @@ class ContextBuilder:
|
||||
entry = entries[index]
|
||||
if (
|
||||
entry.get("session_key") == session_key
|
||||
and entry.get("content") == session_summary.text
|
||||
and entry.get("content") == session_summary["text"]
|
||||
):
|
||||
return [*entries[:index], *entries[index + 1:]]
|
||||
return entries
|
||||
|
||||
+33
-54
@@ -27,7 +27,7 @@ from nanobot.session.manager import (
|
||||
SessionManager,
|
||||
replay_max_messages_for_context,
|
||||
)
|
||||
from nanobot.session.summary import SessionSummary
|
||||
from nanobot.session.summary import session_summary_from_metadata
|
||||
from nanobot.utils.gitstore import GitStore
|
||||
from nanobot.utils.helpers import (
|
||||
content_with_media_breadcrumbs,
|
||||
@@ -38,7 +38,6 @@ from nanobot.utils.helpers import (
|
||||
recent_message_start_index,
|
||||
strip_think,
|
||||
truncate_text,
|
||||
truncate_text_to_tokens,
|
||||
)
|
||||
from nanobot.utils.prompt_templates import render_template
|
||||
from nanobot.utils.workspace_prompts import (
|
||||
@@ -928,10 +927,10 @@ class Consolidator:
|
||||
len(chunk),
|
||||
replay_max_messages,
|
||||
)
|
||||
summary = await self.archive(
|
||||
chunk,
|
||||
summary = await self.archive_session(
|
||||
session,
|
||||
archive_end=end_idx,
|
||||
runtime=runtime,
|
||||
session_key=session.key,
|
||||
)
|
||||
session.last_consolidated = end_idx
|
||||
session.provider_state = None
|
||||
@@ -955,7 +954,7 @@ class Consolidator:
|
||||
"""Estimate prompt size from the full replayable session history."""
|
||||
history = self._full_replay_history(session)
|
||||
channel = session.key.split(":", 1)[0] if ":" in session.key else None
|
||||
summary = SessionSummary.from_metadata(
|
||||
summary = session_summary_from_metadata(
|
||||
session.metadata,
|
||||
fallback_last_active=session.updated_at,
|
||||
)
|
||||
@@ -982,52 +981,24 @@ class Consolidator:
|
||||
- self._SAFETY_BUFFER
|
||||
)
|
||||
|
||||
def _truncate_to_token_budget(self, text: str, *, runtime: LLMRuntime) -> str:
|
||||
"""Truncate text so it fits within the consolidation LLM's token budget."""
|
||||
budget = self._input_token_budget(runtime)
|
||||
if budget <= 0:
|
||||
return truncate_text(text, _RAW_ARCHIVE_MAX_CHARS)
|
||||
return truncate_text_to_tokens(text, budget)
|
||||
|
||||
async def archive(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
*,
|
||||
runtime: LLMRuntime,
|
||||
session_key: str | None = None,
|
||||
summary_messages: list[dict[str, Any]] | None = None,
|
||||
request_messages: list[dict[str, Any]] | None = None,
|
||||
request_tools: list[dict[str, Any]] | None = None,
|
||||
session_key: str,
|
||||
request_messages: list[dict[str, Any]],
|
||||
request_tools: list[dict[str, Any]],
|
||||
) -> str | None:
|
||||
"""Summarize messages and append the result to history.jsonl.
|
||||
|
||||
``summary_messages`` adds context but is excluded from raw fallback.
|
||||
``request_messages`` preserves a prebuilt model-facing prefix instead
|
||||
of flattening the messages; tools are included but disabled.
|
||||
"""
|
||||
"""Execute a prepared consolidation request and persist its result."""
|
||||
if not messages:
|
||||
return None
|
||||
prebuilt_request = request_messages is not None
|
||||
if request_messages is None:
|
||||
formatted = MemoryStore._format_messages(
|
||||
public_history_messages(
|
||||
summary_messages if summary_messages is not None else messages
|
||||
)
|
||||
)
|
||||
formatted = self._truncate_to_token_budget(formatted, runtime=runtime)
|
||||
request_messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": render_template("agent/consolidator_archive.md", strip=True),
|
||||
},
|
||||
{"role": "user", "content": formatted},
|
||||
]
|
||||
try:
|
||||
response = await runtime.provider.chat_with_retry(
|
||||
model=runtime.model,
|
||||
messages=request_messages,
|
||||
tools=request_tools if prebuilt_request else None,
|
||||
tool_choice="none" if prebuilt_request else None,
|
||||
tools=request_tools,
|
||||
tool_choice="none",
|
||||
temperature=runtime.generation.temperature,
|
||||
max_tokens=runtime.generation.max_tokens,
|
||||
reasoning_effort=runtime.generation.reasoning_effort,
|
||||
@@ -1056,23 +1027,31 @@ class Consolidator:
|
||||
)
|
||||
return summary
|
||||
|
||||
async def _archive_idle_tail(
|
||||
async def archive_session(
|
||||
self,
|
||||
session: Session,
|
||||
messages: list[dict[str, Any]],
|
||||
*,
|
||||
archive_end: int,
|
||||
runtime: LLMRuntime,
|
||||
) -> str | None:
|
||||
"""Archive an idle tail by extending the ordinary model-facing messages."""
|
||||
"""Archive a session prefix by appending a consolidation instruction."""
|
||||
messages = list(session.messages[session.last_consolidated:archive_end])
|
||||
if not messages:
|
||||
return None
|
||||
budget = self._input_token_budget(runtime)
|
||||
if budget <= 0:
|
||||
logger.debug(
|
||||
"Idle consolidation has no safe input budget for {}; raw-dumping",
|
||||
"Consolidation has no safe input budget for {}; raw-dumping",
|
||||
session.key,
|
||||
)
|
||||
self.store.raw_archive(messages, session_key=session.key)
|
||||
return None
|
||||
history = session.get_history(
|
||||
prefix = Session(
|
||||
key=session.key,
|
||||
messages=list(session.messages[:archive_end]),
|
||||
last_consolidated=session.last_consolidated,
|
||||
)
|
||||
history = prefix.get_history(
|
||||
max_messages=replay_max_messages_for_context(runtime.context_window_tokens),
|
||||
max_tokens=budget,
|
||||
)
|
||||
@@ -1085,7 +1064,7 @@ class Consolidator:
|
||||
or history[-len(archive_history):] != archive_history
|
||||
):
|
||||
logger.debug(
|
||||
"Idle consolidation cannot replay the full tail for {}; raw-dumping",
|
||||
"Consolidation cannot replay the full chunk for {}; raw-dumping",
|
||||
session.key,
|
||||
)
|
||||
self.store.raw_archive(messages, session_key=session.key)
|
||||
@@ -1103,7 +1082,7 @@ class Consolidator:
|
||||
history=history,
|
||||
current_message=prompt,
|
||||
channel=channel,
|
||||
session_summary=SessionSummary.from_metadata(
|
||||
session_summary=session_summary_from_metadata(
|
||||
session.metadata,
|
||||
fallback_last_active=session.updated_at,
|
||||
),
|
||||
@@ -1120,7 +1099,7 @@ class Consolidator:
|
||||
)
|
||||
if estimated > budget:
|
||||
logger.debug(
|
||||
"Idle consolidation prefix exceeds budget for {}; raw-dumping: {}/{} via {}",
|
||||
"Consolidation prefix exceeds budget for {}; raw-dumping: {}/{} via {}",
|
||||
session.key,
|
||||
estimated,
|
||||
budget,
|
||||
@@ -1215,13 +1194,13 @@ class Consolidator:
|
||||
source,
|
||||
len(chunk),
|
||||
)
|
||||
summary = await self.archive(
|
||||
chunk,
|
||||
summary = await self.archive_session(
|
||||
session,
|
||||
archive_end=end_idx,
|
||||
runtime=runtime,
|
||||
session_key=session.key,
|
||||
)
|
||||
# Advance the cursor either way: on success the chunk was
|
||||
# summarized; on failure archive() already raw-archived it as
|
||||
# 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:
|
||||
@@ -1278,9 +1257,9 @@ class Consolidator:
|
||||
|
||||
last_active = session.updated_at
|
||||
archive_end = archive_start + len(messages_to_archive)
|
||||
summary = await self._archive_idle_tail(
|
||||
summary = await self.archive_session(
|
||||
session,
|
||||
messages_to_archive,
|
||||
archive_end=archive_end,
|
||||
runtime=runtime,
|
||||
)
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ import subprocess
|
||||
import sys
|
||||
import time
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import TYPE_CHECKING, Any, Literal, cast
|
||||
|
||||
from nanobot import __version__
|
||||
@@ -307,18 +307,26 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage:
|
||||
loop.discard_session_file_state(ctx.key)
|
||||
session = ctx.session or loop.sessions.get_or_create(ctx.key)
|
||||
snapshot = session.messages[session.last_consolidated:]
|
||||
archive_snapshot = None
|
||||
runtime = None
|
||||
if snapshot:
|
||||
runtime = ctx.runtime or loop.runtime_for_session(session)
|
||||
archive_snapshot = replace(
|
||||
session,
|
||||
messages=snapshot,
|
||||
metadata=dict(session.metadata),
|
||||
last_consolidated=0,
|
||||
provider_state=None,
|
||||
)
|
||||
session.clear()
|
||||
loop.sessions.save(session)
|
||||
loop.sessions.invalidate(session.key)
|
||||
if snapshot and runtime is not None:
|
||||
if archive_snapshot is not None and runtime is not None:
|
||||
loop.schedule_background(
|
||||
loop.consolidator.archive( # pyright: ignore[reportUnknownMemberType]
|
||||
snapshot,
|
||||
loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType]
|
||||
archive_snapshot,
|
||||
archive_end=len(snapshot),
|
||||
runtime=runtime,
|
||||
session_key=ctx.key,
|
||||
)
|
||||
)
|
||||
return OutboundMessage(
|
||||
|
||||
+24
-35
@@ -1,47 +1,36 @@
|
||||
"""Structured session-summary values used while building model context."""
|
||||
"""Helpers for validated session-summary metadata."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import cast
|
||||
from typing import TypedDict, cast
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SessionSummary:
|
||||
"""A consolidated session checkpoint before presentation formatting."""
|
||||
|
||||
class SessionSummary(TypedDict):
|
||||
text: str
|
||||
last_active: datetime
|
||||
last_active: str
|
||||
|
||||
def for_prompt(self) -> str:
|
||||
return (
|
||||
f"Previous conversation summary (last active {self.last_active.isoformat()}):\n"
|
||||
f"{self.text}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_metadata(
|
||||
cls,
|
||||
metadata: Mapping[str, object] | None,
|
||||
*,
|
||||
fallback_last_active: datetime,
|
||||
) -> SessionSummary | None:
|
||||
raw: object = metadata.get("_last_summary") if metadata is not None else None
|
||||
if not isinstance(raw, Mapping):
|
||||
return None
|
||||
summary_data = cast(Mapping[str, object], raw)
|
||||
text = summary_data.get("text")
|
||||
if not isinstance(text, str) or not text:
|
||||
return None
|
||||
raw_last_active = summary_data.get("last_active")
|
||||
def session_summary_from_metadata(
|
||||
metadata: Mapping[str, object] | None,
|
||||
*,
|
||||
fallback_last_active: datetime,
|
||||
) -> SessionSummary | None:
|
||||
raw: object = metadata.get("_last_summary") if metadata is not None else None
|
||||
if not isinstance(raw, Mapping):
|
||||
return None
|
||||
summary_data = cast(Mapping[str, object], raw)
|
||||
text = summary_data.get("text")
|
||||
if not isinstance(text, str) or not text:
|
||||
return None
|
||||
raw_last_active = summary_data.get("last_active")
|
||||
if isinstance(raw_last_active, str):
|
||||
try:
|
||||
last_active = (
|
||||
datetime.fromisoformat(raw_last_active)
|
||||
if isinstance(raw_last_active, str)
|
||||
else fallback_last_active
|
||||
)
|
||||
datetime.fromisoformat(raw_last_active)
|
||||
last_active = raw_last_active
|
||||
except ValueError:
|
||||
last_active = fallback_last_active
|
||||
return cls(text=text, last_active=last_active)
|
||||
last_active = fallback_last_active.isoformat()
|
||||
else:
|
||||
last_active = fallback_last_active.isoformat()
|
||||
return {"text": text, "last_active": last_active}
|
||||
|
||||
@@ -421,7 +421,7 @@ class TestAutoCompact:
|
||||
|
||||
entry = loop.auto_compact._summaries.get("cli:test")
|
||||
assert entry is not None
|
||||
assert entry.text == "User said hello."
|
||||
assert entry["text"] == "User said hello."
|
||||
session_after = loop.sessions.get_or_create("cli:test")
|
||||
assert len(session_after.messages) == 12
|
||||
assert len(session_after.get_history(max_messages=12)) == (
|
||||
@@ -909,7 +909,7 @@ class TestProactiveAutoCompact:
|
||||
assert len(archived_messages) == 10
|
||||
entry = loop.auto_compact._summaries.get("cli:test")
|
||||
assert entry is not None
|
||||
assert entry.text == "User chatted about old things."
|
||||
assert entry["text"] == "User chatted about old things."
|
||||
await loop.aclose()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -1227,8 +1227,7 @@ class TestSummaryPersistence:
|
||||
_, summary = loop.auto_compact.prepare_session(reloaded, "cli:test")
|
||||
|
||||
assert summary is not None
|
||||
assert summary.text == "User said hello."
|
||||
assert "Previous conversation summary" in summary.for_prompt()
|
||||
assert summary["text"] == "User said hello."
|
||||
# _last_summary persists in metadata for restart survival.
|
||||
assert "_last_summary" in reloaded.metadata
|
||||
await loop.aclose()
|
||||
@@ -1256,7 +1255,7 @@ class TestSummaryPersistence:
|
||||
assert summary is not None
|
||||
_, summary2 = loop.auto_compact.prepare_session(reloaded, "cli:test")
|
||||
assert summary2 is not None
|
||||
assert summary2.text == "Summary."
|
||||
assert summary2["text"] == "Summary."
|
||||
# _last_summary persists in metadata for restart survival.
|
||||
assert "_last_summary" in reloaded.metadata
|
||||
await loop.aclose()
|
||||
@@ -1306,7 +1305,7 @@ class TestSummaryPersistence:
|
||||
loop.sessions.get_or_create("cli:test"), "cli:test"
|
||||
)
|
||||
assert summary1 is not None
|
||||
assert summary1.text == "First summary."
|
||||
assert summary1["text"] == "First summary."
|
||||
assert "cli:test" not in loop.auto_compact._summaries # popped by hot path
|
||||
|
||||
# Add new messages and archive again (simulating a later turn)
|
||||
@@ -1326,7 +1325,7 @@ class TestSummaryPersistence:
|
||||
reloaded = loop.sessions.get_or_create("cli:test")
|
||||
_, summary2 = loop.auto_compact.prepare_session(reloaded, "cli:test")
|
||||
assert summary2 is not None
|
||||
assert summary2.text == "Second summary."
|
||||
assert summary2["text"] == "Second summary."
|
||||
await loop.aclose()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -7,7 +7,6 @@ import pytest
|
||||
|
||||
from nanobot.agent.autocompact import AutoCompact
|
||||
from nanobot.session.manager import Session, SessionManager
|
||||
from nanobot.session.summary import SessionSummary
|
||||
|
||||
|
||||
def _runtime(_session: Session | None = None):
|
||||
@@ -176,24 +175,6 @@ class TestIsExpired:
|
||||
assert ac._is_expired(expired, now=now) is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SessionSummary
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSessionSummary:
|
||||
"""Test prompt rendering for the structured summary value."""
|
||||
|
||||
def test_formats_prompt(self):
|
||||
last_active = datetime(2026, 5, 13, 14, 30, 0)
|
||||
summary = SessionSummary("User discussed Python.", last_active)
|
||||
|
||||
assert summary.for_prompt() == (
|
||||
"Previous conversation summary (last active 2026-05-13T14:30:00):\n"
|
||||
"User discussed Python."
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# check_expired
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -490,7 +471,7 @@ class TestArchiveDelegates:
|
||||
|
||||
entry = ac._summaries.get("cli:test")
|
||||
assert entry is not None
|
||||
assert entry.text == "Hello."
|
||||
assert entry["text"] == "Hello."
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_summary_when_compact_returns_empty(self):
|
||||
@@ -569,21 +550,29 @@ class TestPrepareSession:
|
||||
ac = _make_autocompact()
|
||||
session = _make_session()
|
||||
last_active = datetime(2026, 5, 13, 14, 0, 0)
|
||||
ac._summaries["cli:test"] = SessionSummary("Hot summary.", last_active)
|
||||
ac._summaries["cli:test"] = {
|
||||
"text": "Hot summary.",
|
||||
"last_active": last_active.isoformat(),
|
||||
}
|
||||
|
||||
result_session, summary = ac.prepare_session(session, "cli:test")
|
||||
|
||||
assert result_session is session
|
||||
assert summary is not None
|
||||
assert summary.text == "Hot summary."
|
||||
assert "Previous conversation summary" in summary.for_prompt()
|
||||
assert summary == {
|
||||
"text": "Hot summary.",
|
||||
"last_active": last_active.isoformat(),
|
||||
}
|
||||
|
||||
def test_hot_path_pops_summary_one_shot(self):
|
||||
"""Hot path should pop the summary (one-shot; second call returns None)."""
|
||||
ac = _make_autocompact()
|
||||
session = _make_session()
|
||||
last_active = datetime(2026, 1, 1)
|
||||
ac._summaries["cli:test"] = SessionSummary("One-shot.", last_active)
|
||||
ac._summaries["cli:test"] = {
|
||||
"text": "One-shot.",
|
||||
"last_active": last_active.isoformat(),
|
||||
}
|
||||
|
||||
_, summary1 = ac.prepare_session(session, "cli:test")
|
||||
assert summary1 is not None
|
||||
@@ -606,7 +595,7 @@ class TestPrepareSession:
|
||||
|
||||
assert result_session is session
|
||||
assert summary is not None
|
||||
assert summary.text == "Cold summary."
|
||||
assert summary["text"] == "Cold summary."
|
||||
|
||||
def test_cold_path_tolerates_malformed_last_active(self):
|
||||
"""A malformed persisted last_active must not raise on the turn path.
|
||||
@@ -629,8 +618,10 @@ class TestPrepareSession:
|
||||
|
||||
assert result_session is session
|
||||
assert summary is not None
|
||||
assert summary.text == "Cold summary."
|
||||
assert summary.last_active == fallback
|
||||
assert summary == {
|
||||
"text": "Cold summary.",
|
||||
"last_active": fallback.isoformat(),
|
||||
}
|
||||
|
||||
def test_cold_path_tolerates_missing_last_active(self):
|
||||
"""A _last_summary dict without last_active must not raise."""
|
||||
@@ -645,8 +636,10 @@ class TestPrepareSession:
|
||||
|
||||
assert result_session is session
|
||||
assert summary is not None
|
||||
assert summary.text == "Cold summary."
|
||||
assert summary.last_active == fallback
|
||||
assert summary == {
|
||||
"text": "Cold summary.",
|
||||
"last_active": fallback.isoformat(),
|
||||
}
|
||||
|
||||
def test_cold_path_missing_text_returns_none(self):
|
||||
"""A _last_summary without a non-empty string text yields no summary."""
|
||||
@@ -677,7 +670,10 @@ class TestPrepareSession:
|
||||
ac.sessions = mock_sm
|
||||
key = "dream:20260602-155256"
|
||||
ac._archiving.add(key)
|
||||
ac._summaries[key] = SessionSummary("Hot summary.", datetime(2026, 6, 2, 15, 52, 56))
|
||||
ac._summaries[key] = {
|
||||
"text": "Hot summary.",
|
||||
"last_active": "2026-06-02T15:52:56",
|
||||
}
|
||||
session = _make_session(
|
||||
key=key,
|
||||
updated_at=datetime.now() - timedelta(minutes=20),
|
||||
@@ -717,9 +713,12 @@ class TestPrepareSession:
|
||||
},
|
||||
})
|
||||
last_active = datetime(2026, 5, 13, 14, 0, 0)
|
||||
ac._summaries["cli:test"] = SessionSummary("Hot summary.", last_active)
|
||||
ac._summaries["cli:test"] = {
|
||||
"text": "Hot summary.",
|
||||
"last_active": last_active.isoformat(),
|
||||
}
|
||||
|
||||
_, summary = ac.prepare_session(session, "cli:test")
|
||||
assert summary is not None
|
||||
assert summary.text == "Hot summary."
|
||||
assert summary["text"] == "Hot summary."
|
||||
# After hot path pops, cold path would kick in on next call
|
||||
|
||||
@@ -520,14 +520,14 @@ class TestNewCommandArchival:
|
||||
call_count = 0
|
||||
expected_runtime = loop.llm_runtime()
|
||||
|
||||
async def _failing_summarize(_messages, *, runtime, session_key=None) -> bool:
|
||||
async def _failing_summarize(session, *, archive_end, runtime) -> None:
|
||||
nonlocal call_count
|
||||
assert runtime is expected_runtime
|
||||
assert session_key == "cli:test"
|
||||
assert session.key == "cli:test"
|
||||
assert archive_end == len(session.messages)
|
||||
call_count += 1
|
||||
return False
|
||||
|
||||
loop.consolidator.archive = _failing_summarize # type: ignore[method-assign]
|
||||
loop.consolidator.archive_session = _failing_summarize # type: ignore[method-assign]
|
||||
|
||||
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
|
||||
response = await loop._process_message(new_msg, runtime=expected_runtime)
|
||||
@@ -557,14 +557,14 @@ class TestNewCommandArchival:
|
||||
archived_session_key = None
|
||||
expected_runtime = loop.llm_runtime()
|
||||
|
||||
async def _fake_summarize(messages, *, runtime, session_key=None) -> bool:
|
||||
async def _fake_summarize(session, *, archive_end, runtime) -> str:
|
||||
nonlocal archived_count, archived_session_key
|
||||
assert runtime is expected_runtime
|
||||
archived_count = len(messages)
|
||||
archived_session_key = session_key
|
||||
return True
|
||||
archived_count = len(session.messages[:archive_end])
|
||||
archived_session_key = session.key
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _fake_summarize # type: ignore[method-assign]
|
||||
loop.consolidator.archive_session = _fake_summarize # type: ignore[method-assign]
|
||||
|
||||
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
|
||||
response = await loop._process_message(new_msg, runtime=expected_runtime)
|
||||
@@ -588,12 +588,13 @@ class TestNewCommandArchival:
|
||||
loop.sessions.save(session)
|
||||
expected_runtime = loop.llm_runtime()
|
||||
|
||||
async def _ok_summarize(_messages, *, runtime, session_key=None) -> bool:
|
||||
async def _ok_summarize(session, *, archive_end, runtime) -> str:
|
||||
assert runtime is expected_runtime
|
||||
assert session_key == "cli:test"
|
||||
return True
|
||||
assert session.key == "cli:test"
|
||||
assert archive_end == len(session.messages)
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _ok_summarize # type: ignore[method-assign]
|
||||
loop.consolidator.archive_session = _ok_summarize # type: ignore[method-assign]
|
||||
|
||||
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
|
||||
response = await loop._process_message(new_msg, runtime=expected_runtime)
|
||||
@@ -618,14 +619,15 @@ class TestNewCommandArchival:
|
||||
release_archive = asyncio.Event()
|
||||
expected_runtime = loop.llm_runtime()
|
||||
|
||||
async def _slow_summarize(_messages, *, runtime, session_key=None) -> bool:
|
||||
async def _slow_summarize(session, *, archive_end, runtime) -> str:
|
||||
assert runtime is expected_runtime
|
||||
assert session_key == "cli:test"
|
||||
assert session.key == "cli:test"
|
||||
assert archive_end == len(session.messages)
|
||||
await release_archive.wait()
|
||||
archived.set()
|
||||
return True
|
||||
return "Summary."
|
||||
|
||||
loop.consolidator.archive = _slow_summarize # type: ignore[method-assign]
|
||||
loop.consolidator.archive_session = _slow_summarize # type: ignore[method-assign]
|
||||
|
||||
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
|
||||
await loop._process_message(new_msg, runtime=expected_runtime)
|
||||
|
||||
@@ -72,7 +72,7 @@ async def test_consolidation_ratio_controls_target(
|
||||
context_window_tokens=context_window_tokens,
|
||||
consolidation_ratio=ratio,
|
||||
)
|
||||
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign]
|
||||
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
|
||||
session = _session_with_turns(loop, turns=10)
|
||||
|
||||
remaining_estimates = list(estimates)
|
||||
@@ -90,7 +90,7 @@ async def test_consolidation_ratio_controls_target(
|
||||
runtime=runtime,
|
||||
)
|
||||
|
||||
assert loop.consolidator.archive.await_count == expected_archives
|
||||
assert loop.consolidator.archive_session.await_count == expected_archives
|
||||
|
||||
|
||||
def test_ratio_propagated_from_config_schema() -> None:
|
||||
|
||||
@@ -98,28 +98,20 @@ def _build_test_messages(**kwargs):
|
||||
]
|
||||
|
||||
|
||||
async def _archive(consolidator, messages, runtime, *, session_key="test:session"):
|
||||
return await consolidator.archive(
|
||||
messages,
|
||||
runtime=runtime,
|
||||
session_key=session_key,
|
||||
request_messages=_build_test_messages(
|
||||
history=messages,
|
||||
current_message="consolidate",
|
||||
),
|
||||
request_tools=[],
|
||||
)
|
||||
|
||||
|
||||
class TestConsolidatorSummarize:
|
||||
async def test_archive_prompt_includes_media_breadcrumb(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
path = "/home/user/.nanobot/media/websocket/upload_photo.png"
|
||||
summary = "User uploaded a photo."
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content=summary,
|
||||
finish_reason="stop",
|
||||
)
|
||||
|
||||
result = await consolidator.archive(
|
||||
[{"role": "user", "content": "please inspect this", "media": [path]}],
|
||||
runtime=runtime,
|
||||
)
|
||||
|
||||
prompt = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"]
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert f"[image: {path}]" in prompt
|
||||
assert result == summary
|
||||
assert [entry["content"] for entry in entries] == [summary]
|
||||
|
||||
def test_format_messages_keeps_media_only_user_turn(self):
|
||||
path = "/home/user/.nanobot/media/websocket/clip.mp4"
|
||||
|
||||
@@ -134,31 +126,6 @@ class TestConsolidatorSummarize:
|
||||
|
||||
assert formatted == f"[2026-07-27] USER: [image: {path}]"
|
||||
|
||||
async def test_archive_excludes_model_only_runtime_context(
|
||||
self, consolidator, mock_provider, runtime
|
||||
):
|
||||
content, marker = append_runtime_context(
|
||||
"ship the feature",
|
||||
[RuntimeContextBlock(source="goal", content="host-only goal guidance")],
|
||||
)
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="User wants to ship the feature.",
|
||||
finish_reason="stop",
|
||||
)
|
||||
|
||||
await consolidator.archive(
|
||||
[{
|
||||
"role": "user",
|
||||
"content": content,
|
||||
RUNTIME_CONTEXT_HISTORY_META: marker,
|
||||
}],
|
||||
runtime=runtime,
|
||||
)
|
||||
|
||||
prompt = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"]
|
||||
assert "ship the feature" in prompt
|
||||
assert "host-only goal guidance" not in prompt
|
||||
|
||||
async def test_archive_uses_captured_generation(
|
||||
self, consolidator, mock_provider, runtime
|
||||
):
|
||||
@@ -180,10 +147,7 @@ class TestConsolidatorSummarize:
|
||||
finish_reason="stop",
|
||||
)
|
||||
|
||||
await consolidator.archive(
|
||||
[{"role": "user", "content": "hello"}],
|
||||
runtime=admitted,
|
||||
)
|
||||
await _archive(consolidator, [{"role": "user", "content": "hello"}], admitted)
|
||||
|
||||
call = mock_provider.chat_with_retry.call_args.kwargs
|
||||
assert call["model"] == admitted.model
|
||||
@@ -202,7 +166,7 @@ class TestConsolidatorSummarize:
|
||||
{"role": "user", "content": "fix the auth bug"},
|
||||
{"role": "assistant", "content": "Done, fixed the race condition."},
|
||||
]
|
||||
result = await consolidator.archive(messages, runtime=runtime)
|
||||
result = await _archive(consolidator, messages, runtime)
|
||||
assert result == "User fixed a bug in the auth module."
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert len(entries) == 1
|
||||
@@ -220,9 +184,10 @@ class TestConsolidatorSummarize:
|
||||
)
|
||||
messages = [{"role": "user", "content": "fix the auth bug"}]
|
||||
|
||||
await consolidator.archive(
|
||||
await _archive(
|
||||
consolidator,
|
||||
messages,
|
||||
runtime=runtime,
|
||||
runtime,
|
||||
session_key="telegram:chat-1",
|
||||
)
|
||||
|
||||
@@ -235,7 +200,7 @@ class TestConsolidatorSummarize:
|
||||
"""On LLM failure, raw-dump messages to HISTORY.md."""
|
||||
mock_provider.chat_with_retry.side_effect = Exception("API error")
|
||||
messages = [{"role": "user", "content": "hello"}]
|
||||
result = await consolidator.archive(messages, runtime=runtime)
|
||||
result = await _archive(consolidator, messages, runtime)
|
||||
assert result is None # no summary on raw dump fallback
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert len(entries) == 1
|
||||
@@ -251,9 +216,10 @@ class TestConsolidatorSummarize:
|
||||
mock_provider.chat_with_retry.side_effect = Exception("API error")
|
||||
messages = [{"role": "user", "content": "hello"}]
|
||||
|
||||
await consolidator.archive(
|
||||
await _archive(
|
||||
consolidator,
|
||||
messages,
|
||||
runtime=runtime,
|
||||
runtime,
|
||||
session_key="slack:chat-2",
|
||||
)
|
||||
|
||||
@@ -261,7 +227,7 @@ class TestConsolidatorSummarize:
|
||||
assert entries[0]["session_key"] == "slack:chat-2"
|
||||
|
||||
async def test_summarize_skips_empty_messages(self, consolidator, runtime):
|
||||
result = await consolidator.archive([], runtime=runtime)
|
||||
result = await _archive(consolidator, [], runtime)
|
||||
assert result is None
|
||||
|
||||
|
||||
@@ -293,7 +259,7 @@ class TestConsolidatorArchiveErrorHandling:
|
||||
{"role": "user", "content": "fix the auth bug"},
|
||||
{"role": "assistant", "content": "Done, fixed the race condition."},
|
||||
]
|
||||
result = await consolidator.archive(messages, runtime=runtime)
|
||||
result = await _archive(consolidator, messages, runtime)
|
||||
assert result is None
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert len(entries) == 1
|
||||
@@ -312,7 +278,7 @@ class TestConsolidatorArchiveErrorHandling:
|
||||
{"role": "user", "content": "fix the auth bug"},
|
||||
{"role": "assistant", "content": "Done."},
|
||||
]
|
||||
result = await consolidator.archive(messages, runtime=runtime)
|
||||
result = await _archive(consolidator, messages, runtime)
|
||||
assert result == "User fixed a bug in the auth module."
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert len(entries) == 1
|
||||
@@ -329,9 +295,10 @@ class TestConsolidatorArchiveErrorHandling:
|
||||
consolidator.store.raw_archive = MagicMock()
|
||||
|
||||
with pytest.raises(OSError, match="disk full"):
|
||||
await consolidator.archive(
|
||||
await _archive(
|
||||
consolidator,
|
||||
[{"role": "user", "content": "important"}],
|
||||
runtime=runtime,
|
||||
runtime,
|
||||
)
|
||||
|
||||
consolidator.store.raw_archive.assert_not_called()
|
||||
@@ -339,15 +306,19 @@ class TestConsolidatorArchiveErrorHandling:
|
||||
async def test_archive_propagates_template_failure_without_raw_archive(
|
||||
self, consolidator, mock_provider, runtime, monkeypatch
|
||||
):
|
||||
runtime = replace(runtime, context_window_tokens=128_000)
|
||||
consolidator.store.raw_archive = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"nanobot.agent.memory.render_template",
|
||||
MagicMock(side_effect=RuntimeError("template failed")),
|
||||
)
|
||||
session = Session(key="test:template")
|
||||
session.add_message("user", "important")
|
||||
|
||||
with pytest.raises(RuntimeError, match="template failed"):
|
||||
await consolidator.archive(
|
||||
[{"role": "user", "content": "important"}],
|
||||
await consolidator.archive_session(
|
||||
session,
|
||||
archive_end=len(session.messages),
|
||||
runtime=runtime,
|
||||
)
|
||||
|
||||
@@ -366,9 +337,9 @@ class TestConsolidatorTokenBudget:
|
||||
session.key = "test:key"
|
||||
consolidator.sessions._session_cache[session.key] = session
|
||||
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken"))
|
||||
consolidator.archive = AsyncMock(return_value=True)
|
||||
consolidator.archive_session = AsyncMock(return_value=True)
|
||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||
consolidator.archive.assert_not_called()
|
||||
consolidator.archive_session.assert_not_called()
|
||||
|
||||
async def test_token_estimation_failure_propagates(self, consolidator, runtime):
|
||||
session = Session(key="test:estimate-failure")
|
||||
@@ -434,7 +405,7 @@ class TestConsolidatorTokenBudget:
|
||||
|
||||
consolidator.sessions._session_cache[session.key] = session
|
||||
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken"))
|
||||
consolidator.archive = AsyncMock(return_value="old conversation summary")
|
||||
consolidator.archive_session = AsyncMock(return_value="old conversation summary")
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(
|
||||
session,
|
||||
@@ -442,7 +413,8 @@ class TestConsolidatorTokenBudget:
|
||||
replay_max_messages=6,
|
||||
)
|
||||
|
||||
archived_chunk = consolidator.archive.await_args.args[0]
|
||||
archive_end = consolidator.archive_session.await_args.kwargs["archive_end"]
|
||||
archived_chunk = session.messages[:archive_end]
|
||||
assert archived_chunk[0]["content"] == "u0"
|
||||
assert archived_chunk[-1]["content"] == "a6"
|
||||
assert session.last_consolidated == 14
|
||||
@@ -466,7 +438,7 @@ class TestConsolidatorTokenBudget:
|
||||
|
||||
consolidator.sessions._session_cache[session.key] = session
|
||||
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken"))
|
||||
consolidator.archive = AsyncMock(return_value="tool turn summary")
|
||||
consolidator.archive_session = AsyncMock(return_value="tool turn summary")
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(
|
||||
session,
|
||||
@@ -474,7 +446,8 @@ class TestConsolidatorTokenBudget:
|
||||
replay_max_messages=4,
|
||||
)
|
||||
|
||||
archived_chunk = consolidator.archive.await_args.args[0]
|
||||
archive_end = consolidator.archive_session.await_args.kwargs["archive_end"]
|
||||
archived_chunk = session.messages[:archive_end]
|
||||
assert [m["content"] for m in archived_chunk] == ["old", "old answer"]
|
||||
assert session.last_consolidated == 2
|
||||
|
||||
@@ -501,7 +474,7 @@ class TestConsolidatorTokenBudget:
|
||||
|
||||
consolidator.sessions._session_cache[session.key] = session
|
||||
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken"))
|
||||
consolidator.archive = AsyncMock(return_value="older turn summary")
|
||||
consolidator.archive_session = AsyncMock(return_value="older turn summary")
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(
|
||||
session,
|
||||
@@ -509,7 +482,8 @@ class TestConsolidatorTokenBudget:
|
||||
replay_max_messages=6,
|
||||
)
|
||||
|
||||
archived_chunk = consolidator.archive.await_args.args[0]
|
||||
archive_end = consolidator.archive_session.await_args.kwargs["archive_end"]
|
||||
archived_chunk = session.messages[:archive_end]
|
||||
assert archived_chunk[2]["content"] == "long older turn"
|
||||
assert archived_chunk[-1]["content"] == "older final"
|
||||
assert session.last_consolidated == len(session.messages) - 2
|
||||
@@ -517,12 +491,14 @@ class TestConsolidatorTokenBudget:
|
||||
history = session.get_history(max_messages=6, extend_to_user=True)
|
||||
assert [m["content"] for m in history] == ["new question", "new answer"]
|
||||
|
||||
async def test_large_chunk_archived_without_cap(self, consolidator, runtime):
|
||||
"""Without chunk cap, the full range from pick_consolidation_boundary is archived."""
|
||||
async def test_token_overflow_appends_prompt_to_replay_prefix(
|
||||
self,
|
||||
consolidator,
|
||||
mock_provider,
|
||||
runtime,
|
||||
):
|
||||
consolidator._SAFETY_BUFFER = 0
|
||||
session = MagicMock()
|
||||
session.last_consolidated = 0
|
||||
session.key = "test:key"
|
||||
session = Session(key="test:token-prefix")
|
||||
session.provider_state = _provider_state()
|
||||
session.messages = [
|
||||
{
|
||||
@@ -535,16 +511,24 @@ class TestConsolidatorTokenBudget:
|
||||
consolidator.estimate_session_prompt_tokens = MagicMock(
|
||||
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
|
||||
)
|
||||
# Use real pick_consolidation_boundary — it will find boundary at idx=50
|
||||
# (user message at 50, token budget met)
|
||||
consolidator.archive = AsyncMock(return_value=True)
|
||||
consolidator.pick_consolidation_boundary = MagicMock(return_value=(50, 800))
|
||||
consolidator._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(
|
||||
content="Token overflow summary.",
|
||||
finish_reason="stop",
|
||||
)
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||
|
||||
archived_chunk = consolidator.archive.await_args.args[0]
|
||||
# pick_consolidation_boundary returns (50, tokens) — user turn at idx 50
|
||||
assert archived_chunk[0]["content"] == "m0"
|
||||
assert session.last_consolidated > 0
|
||||
request = mock_provider.chat_with_retry.await_args.kwargs
|
||||
assert [message["content"] for message in request["messages"][1:-1]] == [
|
||||
f"m{i}" for i in range(50)
|
||||
]
|
||||
assert "final 50 conversation messages" in request["messages"][-1]["content"]
|
||||
assert request["tools"] == []
|
||||
assert request["tool_choice"] == "none"
|
||||
assert session.last_consolidated == 50
|
||||
assert session.provider_state is None
|
||||
|
||||
async def test_raw_archive_fallback_advances_last_consolidated(
|
||||
@@ -567,12 +551,12 @@ class TestConsolidatorTokenBudget:
|
||||
consolidator.estimate_session_prompt_tokens = MagicMock(
|
||||
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
|
||||
)
|
||||
# LLM consolidation fails — archive() returns None (raw_archive fired).
|
||||
consolidator.archive = AsyncMock(return_value=None)
|
||||
# LLM consolidation fails after raw_archive fires.
|
||||
consolidator.archive_session = AsyncMock(return_value=None)
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||
|
||||
consolidator.archive.assert_awaited_once()
|
||||
consolidator.archive_session.assert_awaited_once()
|
||||
# The chunk is considered "materialized" (as a raw-archive breadcrumb),
|
||||
# so last_consolidated must have moved past it.
|
||||
assert session.last_consolidated == 50
|
||||
@@ -596,12 +580,12 @@ class TestConsolidatorTokenBudget:
|
||||
consolidator.estimate_session_prompt_tokens = MagicMock(
|
||||
return_value=(1200, "tiktoken")
|
||||
)
|
||||
consolidator.archive = AsyncMock(return_value=None)
|
||||
consolidator.archive_session = AsyncMock(return_value=None)
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||
|
||||
# Exactly one fallback per call — not _MAX_CONSOLIDATION_ROUNDS.
|
||||
assert consolidator.archive.await_count == 1
|
||||
assert consolidator.archive_session.await_count == 1
|
||||
|
||||
async def test_boundary_respected_when_no_intermediate_user_turn(
|
||||
self, consolidator, runtime
|
||||
@@ -622,11 +606,11 @@ class TestConsolidatorTokenBudget:
|
||||
consolidator.estimate_session_prompt_tokens = MagicMock(
|
||||
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
|
||||
)
|
||||
consolidator.archive = AsyncMock(return_value=True)
|
||||
consolidator.archive_session = AsyncMock(return_value=True)
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||
|
||||
consolidator.archive.assert_awaited_once()
|
||||
consolidator.archive_session.assert_awaited_once()
|
||||
# pick_consolidation_boundary finds the only boundary at idx=61
|
||||
assert session.last_consolidated == 61
|
||||
|
||||
@@ -1439,43 +1423,7 @@ class TestRawArchiveTruncation:
|
||||
assert len(entries[0]["content"]) < 200
|
||||
|
||||
|
||||
class TestArchiveTruncation:
|
||||
"""archive() must truncate formatted text before sending to consolidation LLM."""
|
||||
|
||||
async def test_archive_truncates_large_formatted_text(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""Large formatted text should be truncated to token budget before LLM call."""
|
||||
# context_window_tokens=1000, max_completion_tokens=100, _SAFETY_BUFFER=1024
|
||||
# budget = 1000 - 100 - 1024 = -124 → fallback via truncate_text(budget*4)
|
||||
big_messages = [{"role": "user", "content": "x" * 100_000}]
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="Summary of large input.", finish_reason="stop"
|
||||
)
|
||||
await consolidator.archive(big_messages, runtime=runtime)
|
||||
|
||||
call_args = mock_provider.chat_with_retry.call_args
|
||||
user_content = call_args.kwargs["messages"][1]["content"]
|
||||
# Should be significantly shorter than 100K
|
||||
assert len(user_content) < 50_000
|
||||
|
||||
async def test_archive_truncates_with_small_token_budget(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""Small context window: truncation uses actual tokenizer count."""
|
||||
runtime = replace(runtime, context_window_tokens=500)
|
||||
big_messages = [{"role": "user", "content": "word " * 50_000}]
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="Summary.", finish_reason="stop"
|
||||
)
|
||||
await consolidator.archive(big_messages, runtime=runtime)
|
||||
|
||||
sent_messages = mock_provider.chat_with_retry.call_args.kwargs["messages"]
|
||||
user_content = sent_messages[1]["content"]
|
||||
# budget = 500 - 100 - 1024 = negative, fallback char-based
|
||||
# Should be truncated
|
||||
assert len(user_content) < 250_000
|
||||
|
||||
class TestArchivePersistence:
|
||||
async def test_oversized_summary_is_capped_before_append(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
@@ -1486,29 +1434,11 @@ class TestArchiveTruncation:
|
||||
content="S" * (_ARCHIVE_SUMMARY_MAX_CHARS * 10),
|
||||
finish_reason="stop",
|
||||
)
|
||||
await consolidator.archive(
|
||||
await _archive(
|
||||
consolidator,
|
||||
[{"role": "user", "content": "hi"}],
|
||||
runtime=runtime,
|
||||
runtime,
|
||||
)
|
||||
|
||||
entry = store.read_unprocessed_history(since_cursor=0)[0]
|
||||
assert len(entry["content"]) <= _ARCHIVE_SUMMARY_MAX_CHARS + 50
|
||||
|
||||
async def test_archive_truncates_via_tiktoken_with_positive_budget(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""Positive token budget should use tiktoken for precise truncation."""
|
||||
runtime = replace(runtime, context_window_tokens=10_000)
|
||||
consolidator._SAFETY_BUFFER = 0
|
||||
# budget = 10000 - 100 - 0 = 9900 tokens
|
||||
big_messages = [{"role": "user", "content": "word " * 50_000}]
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="Summary.", finish_reason="stop"
|
||||
)
|
||||
await consolidator.archive(big_messages, runtime=runtime)
|
||||
|
||||
import tiktoken
|
||||
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
|
||||
|
||||
@@ -1,13 +1,11 @@
|
||||
"""Tests for ContextBuilder — system prompt and message assembly."""
|
||||
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.context import ContextBuilder
|
||||
from nanobot.runtime_context import RuntimeContextBlock
|
||||
from nanobot.session.summary import SessionSummary
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
@@ -332,7 +330,10 @@ class TestBuildSystemPrompt:
|
||||
|
||||
def test_includes_session_summary(self, tmp_path):
|
||||
builder = _builder(tmp_path)
|
||||
summary = SessionSummary("Previous chat about Python.", datetime(2026, 8, 19, 10, 0))
|
||||
summary = {
|
||||
"text": "Previous chat about Python.",
|
||||
"last_active": "2026-08-19T10:00:00",
|
||||
}
|
||||
result = builder.build_system_prompt(session_summary=summary)
|
||||
assert "Previous chat about Python." in result
|
||||
assert "[Archived Context Summary]" in result
|
||||
@@ -340,7 +341,7 @@ class TestBuildSystemPrompt:
|
||||
def test_sections_separated_by_separator(self, tmp_path):
|
||||
(tmp_path / "AGENTS.md").write_text("Rules.", encoding="utf-8")
|
||||
builder = _builder(tmp_path)
|
||||
summary = SessionSummary("Summary.", datetime(2026, 8, 19, 10, 0))
|
||||
summary = {"text": "Summary.", "last_active": "2026-08-19T10:00:00"}
|
||||
result = builder.build_system_prompt(session_summary=summary)
|
||||
assert "\n\n---\n\n" in result
|
||||
|
||||
|
||||
@@ -10,7 +10,6 @@ from pathlib import Path
|
||||
|
||||
from nanobot.agent.context import ContextBuilder
|
||||
from nanobot.runtime_context import RuntimeContextBlock
|
||||
from nanobot.session.summary import SessionSummary
|
||||
|
||||
|
||||
class _FakeDatetime(real_datetime):
|
||||
@@ -125,7 +124,7 @@ def test_session_summary_replaces_interleaved_recent_history_entry(tmp_path) ->
|
||||
"later telegram event",
|
||||
session_key="telegram:chat-1",
|
||||
)
|
||||
summary = SessionSummary(overview, real_datetime(2026, 8, 19, 10, 0))
|
||||
summary = {"text": overview, "last_active": "2026-08-19T10:00:00"}
|
||||
|
||||
prompt = builder.build_system_prompt(
|
||||
session_key=session_key,
|
||||
|
||||
@@ -7,7 +7,6 @@ from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.providers.base import LLMResponse
|
||||
from nanobot.session.manager import replay_max_messages_for_context
|
||||
from nanobot.session.summary import SessionSummary
|
||||
|
||||
|
||||
def _make_loop(tmp_path, *, estimated_tokens: int, context_window_tokens: int) -> AgentLoop:
|
||||
@@ -35,17 +34,17 @@ def _make_loop(tmp_path, *, estimated_tokens: int, context_window_tokens: int) -
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_below_threshold_does_not_consolidate(tmp_path) -> None:
|
||||
loop = _make_loop(tmp_path, estimated_tokens=100, context_window_tokens=200)
|
||||
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign]
|
||||
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
|
||||
|
||||
await loop.process_direct("hello", session_key="cli:test")
|
||||
|
||||
loop.consolidator.archive.assert_not_awaited()
|
||||
loop.consolidator.archive_session.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypatch) -> None:
|
||||
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
|
||||
loop.consolidator.archive = 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.messages = [
|
||||
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
|
||||
@@ -57,13 +56,13 @@ async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypat
|
||||
|
||||
await loop.process_direct("hello", session_key="cli:test")
|
||||
|
||||
assert loop.consolidator.archive.await_count >= 1
|
||||
assert loop.consolidator.archive_session.await_count >= 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path, monkeypatch) -> None:
|
||||
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
|
||||
loop.consolidator.archive = 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.messages = [
|
||||
@@ -83,7 +82,8 @@ async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path
|
||||
runtime=loop.llm_runtime(),
|
||||
)
|
||||
|
||||
archived_chunk = loop.consolidator.archive.await_args.args[0]
|
||||
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_consolidated == 4
|
||||
|
||||
@@ -92,7 +92,7 @@ async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path
|
||||
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 = 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.messages = [
|
||||
@@ -123,7 +123,7 @@ async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> No
|
||||
runtime=loop.llm_runtime(),
|
||||
)
|
||||
|
||||
assert loop.consolidator.archive.await_count == 2
|
||||
assert loop.consolidator.archive_session.await_count == 2
|
||||
assert session.last_consolidated == 6
|
||||
|
||||
|
||||
@@ -131,7 +131,7 @@ async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> No
|
||||
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 = 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.messages = [
|
||||
@@ -163,14 +163,14 @@ async def test_consolidation_continues_below_trigger_until_half_target(tmp_path,
|
||||
runtime=loop.llm_runtime(),
|
||||
)
|
||||
|
||||
assert loop.consolidator.archive.await_count == 2
|
||||
assert loop.consolidator.archive_session.await_count == 2
|
||||
assert session.last_consolidated == 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.consolidator.archive = 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.messages = [
|
||||
@@ -203,7 +203,7 @@ async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path,
|
||||
|
||||
reloaded, pending = loop.auto_compact.prepare_session(reloaded, "cli:test")
|
||||
assert pending is not None
|
||||
assert pending.text == "User discussed project status."
|
||||
assert pending["text"] == "User discussed project status."
|
||||
# _last_summary persists for restart survival.
|
||||
assert "_last_summary" in reloaded.metadata
|
||||
|
||||
@@ -213,7 +213,10 @@ async def test_preflight_consolidation_receives_pending_summary(tmp_path) -> Non
|
||||
loop = _make_loop(tmp_path, estimated_tokens=100, context_window_tokens=200)
|
||||
session = loop.sessions.get_or_create("cli:test")
|
||||
loop.auto_compact.prepare_session = MagicMock(
|
||||
return_value=(session, SessionSummary("earlier context", session.updated_at))
|
||||
return_value=(
|
||||
session,
|
||||
{"text": "earlier context", "last_active": session.updated_at.isoformat()},
|
||||
)
|
||||
) # type: ignore[method-assign]
|
||||
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None) # type: ignore[method-assign]
|
||||
loop.schedule_background = lambda coro: coro.close() # type: ignore[method-assign]
|
||||
@@ -242,11 +245,11 @@ async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) ->
|
||||
|
||||
archived_session_keys: list[str | None] = []
|
||||
|
||||
async def track_consolidate(messages, *, runtime, session_key=None):
|
||||
async def track_consolidate(session, *, archive_end, runtime):
|
||||
order.append("consolidate")
|
||||
archived_session_keys.append(session_key)
|
||||
archived_session_keys.append(session.key)
|
||||
return True
|
||||
loop.consolidator.archive = track_consolidate # type: ignore[method-assign]
|
||||
loop.consolidator.archive_session = track_consolidate # type: ignore[method-assign]
|
||||
|
||||
async def track_llm(*args, **kwargs):
|
||||
order.append("llm")
|
||||
|
||||
@@ -258,7 +258,7 @@ class TestCmdNewUnifiedSession:
|
||||
previous_file_state.record_read(tracked_file)
|
||||
loop = SimpleNamespace(
|
||||
sessions=sessions,
|
||||
consolidator=SimpleNamespace(archive=AsyncMock(return_value=True)),
|
||||
consolidator=SimpleNamespace(archive_session=AsyncMock(return_value=True)),
|
||||
_cancel_active_tasks=AsyncMock(return_value=0),
|
||||
discard_session_file_state=file_state_store.discard,
|
||||
llm_runtime=MagicMock(return_value=MagicMock()),
|
||||
@@ -288,10 +288,14 @@ class TestCmdNewUnifiedSession:
|
||||
reset_file_state = file_state_store.for_session("unified:default")
|
||||
assert reset_file_state is not previous_file_state
|
||||
assert reset_file_state.is_unchanged(tracked_file) is False
|
||||
loop.consolidator.archive.assert_called_once_with(
|
||||
expected_snapshot,
|
||||
archived = loop.consolidator.archive_session.call_args.args[0]
|
||||
assert archived.key == "unified:default"
|
||||
assert archived.messages == expected_snapshot
|
||||
assert archived.last_consolidated == 0
|
||||
loop.consolidator.archive_session.assert_called_once_with(
|
||||
archived,
|
||||
archive_end=len(expected_snapshot),
|
||||
runtime=admitted_runtime,
|
||||
session_key="unified:default",
|
||||
)
|
||||
loop.llm_runtime.assert_not_called()
|
||||
|
||||
@@ -310,7 +314,7 @@ class TestCmdNewUnifiedSession:
|
||||
|
||||
loop = SimpleNamespace(
|
||||
sessions=sessions,
|
||||
consolidator=SimpleNamespace(archive=AsyncMock(return_value=True)),
|
||||
consolidator=SimpleNamespace(archive_session=AsyncMock(return_value=True)),
|
||||
_cancel_active_tasks=AsyncMock(return_value=0),
|
||||
discard_session_file_state=MagicMock(),
|
||||
runtime_for_session=MagicMock(return_value=MagicMock()),
|
||||
@@ -356,7 +360,7 @@ class TestConsolidationUnaffectedByUnifiedSession:
|
||||
build_messages=MagicMock(return_value=[]),
|
||||
get_tool_definitions=MagicMock(return_value=[]),
|
||||
)
|
||||
consolidator.archive = AsyncMock()
|
||||
consolidator.archive_session = AsyncMock()
|
||||
|
||||
session = Session(key="unified:default")
|
||||
session.messages = []
|
||||
@@ -364,11 +368,11 @@ class TestConsolidationUnaffectedByUnifiedSession:
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||
|
||||
consolidator.archive.assert_not_called()
|
||||
consolidator.archive_session.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consolidation_behaviour_identical_for_any_key(self):
|
||||
"""archive call count is the same for 'telegram:123' and 'unified:default'
|
||||
"""Archive call count is the same for 'telegram:123' and 'unified:default'
|
||||
under identical token conditions."""
|
||||
from nanobot.agent.memory import Consolidator, MemoryStore
|
||||
|
||||
@@ -392,12 +396,12 @@ class TestConsolidationUnaffectedByUnifiedSession:
|
||||
session.messages = [] # empty → exits immediately for both keys
|
||||
sessions.get_or_create.return_value = session
|
||||
|
||||
consolidator.archive = AsyncMock()
|
||||
consolidator.archive_session = AsyncMock()
|
||||
await consolidator.maybe_consolidate_by_tokens(
|
||||
session,
|
||||
runtime=runtime,
|
||||
)
|
||||
archive_calls[key] = consolidator.archive.call_count
|
||||
archive_calls[key] = consolidator.archive_session.call_count
|
||||
|
||||
assert archive_calls["telegram:123"] == archive_calls["unified:default"] == 0
|
||||
|
||||
@@ -427,7 +431,7 @@ class TestConsolidationUnaffectedByUnifiedSession:
|
||||
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(950, "tiktoken"))
|
||||
# No valid boundary found → returns gracefully without archiving
|
||||
consolidator.pick_consolidation_boundary = MagicMock(return_value=None)
|
||||
consolidator.archive = AsyncMock()
|
||||
consolidator.archive_session = AsyncMock()
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||
|
||||
@@ -437,7 +441,7 @@ class TestConsolidationUnaffectedByUnifiedSession:
|
||||
runtime=runtime,
|
||||
)
|
||||
# but archive was not called (no valid boundary)
|
||||
consolidator.archive.assert_not_called()
|
||||
consolidator.archive_session.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user