| 1 | """Auto compact: proactive compression of idle sessions to reduce token cost and latency.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from collections.abc import Collection |
| 6 | from datetime import datetime |
| 7 | from typing import TYPE_CHECKING, Any, Callable, Coroutine |
| 8 | |
| 9 | from loguru import logger |
| 10 | from nanobot.session.manager import Session, SessionManager |
| 11 | |
| 12 | if TYPE_CHECKING: |
| 13 | from nanobot.agent.memory import Consolidator |
| 14 | |
| 15 | |
| 16 | class AutoCompact: |
| 17 | _RECENT_SUFFIX_MESSAGES = 8 |
| 18 | |
| 19 | def __init__(self, sessions: SessionManager, consolidator: Consolidator, |
| 20 | session_ttl_minutes: int = 0): |
| 21 | self.sessions = sessions |
| 22 | self.consolidator = consolidator |
| 23 | self._ttl = session_ttl_minutes |
| 24 | self._archiving: set[str] = set() |
| 25 | self._summaries: dict[str, tuple[str, datetime]] = {} |
| 26 | |
| 27 | def _is_expired(self, ts: datetime | str | None, |
| 28 | now: datetime | None = None) -> bool: |
| 29 | if self._ttl <= 0 or not ts: |
| 30 | return False |
| 31 | if isinstance(ts, str): |
| 32 | ts = datetime.fromisoformat(ts) |
| 33 | return ((now or datetime.now()) - ts).total_seconds() >= self._ttl * 60 |
| 34 | |
| 35 | @staticmethod |
| 36 | def _format_summary(text: str, last_active: datetime) -> str: |
| 37 | idle_min = int((datetime.now() - last_active).total_seconds() / 60) |
| 38 | return f"Inactive for {idle_min} minutes.\nPrevious conversation summary: {text}" |
| 39 | |
| 40 | def _split_unconsolidated( |
| 41 | self, session: Session, |
| 42 | ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: |
| 43 | """Split live session tail into archiveable prefix and retained recent suffix.""" |
| 44 | tail = list(session.messages[session.last_consolidated:]) |
| 45 | if not tail: |
| 46 | return [], [] |
| 47 | |
| 48 | probe = Session( |
| 49 | key=session.key, |
| 50 | messages=tail.copy(), |
| 51 | created_at=session.created_at, |
| 52 | updated_at=session.updated_at, |
| 53 | metadata={}, |
| 54 | last_consolidated=0, |
| 55 | ) |
| 56 | probe.retain_recent_legal_suffix(self._RECENT_SUFFIX_MESSAGES) |
| 57 | kept = probe.messages |
| 58 | cut = len(tail) - len(kept) |
| 59 | return tail[:cut], kept |
| 60 | |
| 61 | def check_expired(self, schedule_background: Callable[[Coroutine], None], |
| 62 | active_session_keys: Collection[str] = ()) -> None: |
| 63 | """Schedule archival for idle sessions, skipping those with in-flight agent tasks.""" |
| 64 | now = datetime.now() |
| 65 | for info in self.sessions.list_sessions(): |
| 66 | key = info.get("key", "") |
| 67 | if not key or key in self._archiving: |
| 68 | continue |
| 69 | if key in active_session_keys: |
| 70 | continue |
| 71 | if self._is_expired(info.get("updated_at"), now): |
| 72 | self._archiving.add(key) |
| 73 | schedule_background(self._archive(key)) |
| 74 | |
| 75 | async def _archive(self, key: str) -> None: |
| 76 | try: |
| 77 | self.sessions.invalidate(key) |
| 78 | session = self.sessions.get_or_create(key) |
| 79 | archive_msgs, kept_msgs = self._split_unconsolidated(session) |
| 80 | if not archive_msgs and not kept_msgs: |
| 81 | session.updated_at = datetime.now() |
| 82 | self.sessions.save(session) |
| 83 | return |
| 84 | |
| 85 | last_active = session.updated_at |
| 86 | summary = "" |
| 87 | if archive_msgs: |
| 88 | summary = await self.consolidator.archive(archive_msgs) or "" |
| 89 | if summary and summary != "(nothing)": |
| 90 | self._summaries[key] = (summary, last_active) |
| 91 | session.metadata["_last_summary"] = {"text": summary, "last_active": last_active.isoformat()} |
| 92 | session.messages = kept_msgs |
| 93 | session.last_consolidated = 0 |
| 94 | session.updated_at = datetime.now() |
| 95 | self.sessions.save(session) |
| 96 | if archive_msgs: |
| 97 | logger.info( |
| 98 | "Auto-compact: archived {} (archived={}, kept={}, summary={})", |
| 99 | key, |
| 100 | len(archive_msgs), |
| 101 | len(kept_msgs), |
| 102 | bool(summary), |
| 103 | ) |
| 104 | except Exception: |
| 105 | logger.exception("Auto-compact: failed for {}", key) |
| 106 | finally: |
| 107 | self._archiving.discard(key) |
| 108 | |
| 109 | def prepare_session(self, session: Session, key: str) -> tuple[Session, str | None]: |
| 110 | if key in self._archiving or self._is_expired(session.updated_at): |
| 111 | logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving) |
| 112 | session = self.sessions.get_or_create(key) |
| 113 | # Hot path: summary from in-memory dict (process hasn't restarted). |
| 114 | # Also clean metadata copy so stale _last_summary never leaks to disk. |
| 115 | entry = self._summaries.pop(key, None) |
| 116 | if entry: |
| 117 | session.metadata.pop("_last_summary", None) |
| 118 | return session, self._format_summary(entry[0], entry[1]) |
| 119 | if "_last_summary" in session.metadata: |
| 120 | meta = session.metadata.pop("_last_summary") |
| 121 | self.sessions.save(session) |
| 122 | return session, self._format_summary(meta["text"], datetime.fromisoformat(meta["last_active"])) |
| 123 | return session, None |
| 124 |