chengyongruandXubin Ren
e5a1416a37
The previous implementation popped _last_summary from session.metadata after injecting it into the prompt, then saved the session. This caused the summary to be permanently lost after a process restart, making the AI forget archived context and appear to ignore memory or reference non-existent previous messages. Replace the destructive pop with a _last_summary_used sentinel: - _last_summary stays in metadata for restart survival - _last_summary_used prevents duplicate injection within the same turn - Clear the sentinel whenever a new summary is generated Updates tests to match the new persistence behavior.
127 lines
5.3 KiB
Python
127 lines
5.3 KiB
Python
"""Auto compact: proactive compression of idle sessions to reduce token cost and latency."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Collection
|
|
from datetime import datetime
|
|
from typing import TYPE_CHECKING, Any, Callable, Coroutine
|
|
|
|
from loguru import logger
|
|
from nanobot.session.manager import Session, SessionManager
|
|
|
|
if TYPE_CHECKING:
|
|
from nanobot.agent.memory import Consolidator
|
|
|
|
|
|
class AutoCompact:
|
|
_RECENT_SUFFIX_MESSAGES = 8
|
|
|
|
def __init__(self, sessions: SessionManager, consolidator: Consolidator,
|
|
session_ttl_minutes: int = 0):
|
|
self.sessions = sessions
|
|
self.consolidator = consolidator
|
|
self._ttl = session_ttl_minutes
|
|
self._archiving: set[str] = set()
|
|
self._summaries: dict[str, tuple[str, datetime]] = {}
|
|
|
|
def _is_expired(self, ts: datetime | str | None,
|
|
now: datetime | None = None) -> bool:
|
|
if self._ttl <= 0 or not ts:
|
|
return False
|
|
if isinstance(ts, str):
|
|
ts = datetime.fromisoformat(ts)
|
|
return ((now or datetime.now()) - ts).total_seconds() >= self._ttl * 60
|
|
|
|
@staticmethod
|
|
def _format_summary(text: str, last_active: datetime) -> str:
|
|
idle_min = int((datetime.now() - last_active).total_seconds() / 60)
|
|
return f"Inactive for {idle_min} minutes.\nPrevious conversation summary: {text}"
|
|
|
|
def _split_unconsolidated(
|
|
self, session: Session,
|
|
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
|
"""Split live session tail into archiveable prefix and retained recent suffix."""
|
|
tail = list(session.messages[session.last_consolidated:])
|
|
if not tail:
|
|
return [], []
|
|
|
|
probe = Session(
|
|
key=session.key,
|
|
messages=tail.copy(),
|
|
created_at=session.created_at,
|
|
updated_at=session.updated_at,
|
|
metadata={},
|
|
last_consolidated=0,
|
|
)
|
|
probe.retain_recent_legal_suffix(self._RECENT_SUFFIX_MESSAGES)
|
|
kept = probe.messages
|
|
cut = len(tail) - len(kept)
|
|
return tail[:cut], kept
|
|
|
|
def check_expired(self, schedule_background: Callable[[Coroutine], None],
|
|
active_session_keys: Collection[str] = ()) -> None:
|
|
"""Schedule archival for idle sessions, skipping those with in-flight agent tasks."""
|
|
now = datetime.now()
|
|
for info in self.sessions.list_sessions():
|
|
key = info.get("key", "")
|
|
if not key or key in self._archiving:
|
|
continue
|
|
if key in active_session_keys:
|
|
continue
|
|
if self._is_expired(info.get("updated_at"), now):
|
|
self._archiving.add(key)
|
|
schedule_background(self._archive(key))
|
|
|
|
async def _archive(self, key: str) -> None:
|
|
try:
|
|
self.sessions.invalidate(key)
|
|
session = self.sessions.get_or_create(key)
|
|
archive_msgs, kept_msgs = self._split_unconsolidated(session)
|
|
if not archive_msgs and not kept_msgs:
|
|
session.updated_at = datetime.now()
|
|
self.sessions.save(session)
|
|
return
|
|
|
|
last_active = session.updated_at
|
|
summary = ""
|
|
if archive_msgs:
|
|
summary = await self.consolidator.archive(archive_msgs) or ""
|
|
if summary and summary != "(nothing)":
|
|
self._summaries[key] = (summary, last_active)
|
|
session.metadata["_last_summary"] = {"text": summary, "last_active": last_active.isoformat()}
|
|
session.metadata.pop("_last_summary_used", None)
|
|
session.messages = kept_msgs
|
|
session.last_consolidated = 0
|
|
session.updated_at = datetime.now()
|
|
self.sessions.save(session)
|
|
if archive_msgs:
|
|
logger.info(
|
|
"Auto-compact: archived {} (archived={}, kept={}, summary={})",
|
|
key,
|
|
len(archive_msgs),
|
|
len(kept_msgs),
|
|
bool(summary),
|
|
)
|
|
except Exception:
|
|
logger.exception("Auto-compact: failed for {}", key)
|
|
finally:
|
|
self._archiving.discard(key)
|
|
|
|
def prepare_session(self, session: Session, key: str) -> tuple[Session, str | None]:
|
|
if key in self._archiving or self._is_expired(session.updated_at):
|
|
logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving)
|
|
session = self.sessions.get_or_create(key)
|
|
# Hot path: summary from in-memory dict (process hasn't restarted).
|
|
entry = self._summaries.pop(key, None)
|
|
if entry:
|
|
session.metadata["_last_summary_used"] = True
|
|
return session, self._format_summary(entry[0], entry[1])
|
|
# Cold path: summary persisted in session metadata (process restarted).
|
|
# Keep it in metadata so it survives restarts; only inject once per
|
|
# turn via the _last_summary_used sentinel.
|
|
meta = session.metadata.get("_last_summary")
|
|
if meta and not session.metadata.get("_last_summary_used"):
|
|
session.metadata["_last_summary_used"] = True
|
|
return session, self._format_summary(meta["text"], datetime.fromisoformat(meta["last_active"]))
|
|
return session, None
|