"""Abstract base class for pluggable context engines. A context engine decides when and how conversation context is compacted near the model's token limit, tracks token usage, and may expose tools. The built-in ContextCompressor is the default; ``context.engine`` in config.yaml selects a plugin engine (``plugins/context_engine//``). One engine is active at a time. Lifecycle: on_session_start() -> per API response update_from_response() -> per turn should_compress() / compress() -> on_session_end() at real session boundaries only (CLI exit, /reset, gateway expiry), never per-turn. """ from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional from agent.redact import redact_sensitive_text MEMORY_CONTEXT_MAX_CHARS = 6_000 _MEMORY_CONTEXT_HEAD_CHARS = 4_000 _MEMORY_CONTEXT_TAIL_CHARS = 1_500 _MEMORY_CONTEXT_TRUNCATION_MARKER = "\n...[memory provider context truncated]...\n" def sanitize_memory_context(memory_context: str) -> str: """Prepare provider context for a context-engine/LLM egress boundary.""" sanitized = redact_sensitive_text( memory_context.strip(), force=True, redact_url_credentials=True, ) if len(sanitized) <= MEMORY_CONTEXT_MAX_CHARS: return sanitized return ( sanitized[:_MEMORY_CONTEXT_HEAD_CHARS] + _MEMORY_CONTEXT_TRUNCATION_MARKER + sanitized[-_MEMORY_CONTEXT_TAIL_CHARS:] ) def automatic_compaction_status_message( engine: Any, *, phase: str, default_message: str, **context: Any, ) -> str | None: """Host-visible status for an automatic compaction event; ``None`` = emit nothing. Engines suppress via ``emit_automatic_compaction_status = False`` or customize via ``get_automatic_compaction_status_message(...)``. """ if not getattr(engine, "emit_automatic_compaction_status", True): return None formatter = getattr(engine, "get_automatic_compaction_status_message", None) if callable(formatter): message = formatter( phase=phase, default_message=default_message, **context, ) else: message = default_message if message is None: return None message = str(message).strip() return message or None class ContextEngine(ABC): """Base class all context engines must implement.""" # -- Identity ---------------------------------------------------------- @property @abstractmethod def name(self) -> str: """Short identifier (e.g. 'compressor', 'lcm').""" # -- Token state: engines MUST maintain these; run_agent.py reads them directly. last_prompt_tokens: int = 0 last_completion_tokens: int = 0 last_total_tokens: int = 0 threshold_tokens: int = 0 context_length: int = 0 compression_count: int = 0 # -- Compaction parameters (read by run_agent.py for preflight). protect_first_n # counts non-system head messages kept verbatim IN ADDITION to the always- # protected system prompt (3 keeps the historical head shape). threshold_percent: float = 0.75 protect_first_n: int = 3 protect_last_n: int = 6 # False keeps successful automatic compaction passes silent (routine # background maintenance); warnings, errors and manual /compress still surface. emit_automatic_compaction_status: bool = True # -- Core interface ---------------------------------------------------- @abstractmethod def update_from_response(self, usage: Dict[str, Any]) -> None: """Update tracked token usage after every LLM call. ``prompt_tokens``/``completion_tokens``/``total_tokens`` are always present; the canonical buckets (``input_tokens``, ``output_tokens``, ``cache_read_tokens``, ``cache_write_tokens``, ``reasoning_tokens``) are optional on older hosts. """ @abstractmethod def should_compress(self, prompt_tokens: int = None) -> bool: """Return True if compaction should fire this turn.""" def should_compress_info(self, prompt_tokens: int = None) -> "tuple[bool, str | None]": """Return ``(should_compress, reason)``. Engines with block reasons (summary-LLM cooldown, anti-thrashing guard) override this so callers can warn the user instead of silently skipping compression. The default keeps plugin engines from raising AttributeError. """ return self.should_compress(prompt_tokens), None @abstractmethod def compress( self, messages: List[Dict[str, Any]], current_tokens: Optional[int] = None, focus_topic: Optional[str] = None, force: bool = False, memory_context: str = "", ) -> List[Dict[str, Any]]: """Compact ``messages`` and return a valid OpenAI-format message list that fits the context budget (summarize, build a DAG, anything). ``focus_topic`` comes from manual ``/compress `` (prioritise that topic); ``force`` asks to bypass an engine-owned cooldown; ``memory_context`` is provider text to include in the handoff prompt. Older engines may omit optional parameters — the host filters them by signature. """ # -- Optional: proactive tool-result prune ----------------------------- def prune_tool_results_only( self, messages: List[Dict[str, Any]], current_tokens: int | None = None, ) -> tuple[List[Dict[str, Any]], int]: """Deterministically trim old tool-result payloads without an LLM call. Runs on a low, cost-oriented trigger independent of ``should_compress`` so large-window engines reclaim re-sent tool output long before full compaction. Returns ``(messages, n_pruned)``; default is a no-op so engines predating this hook never raise in the post-tool-call prune path. """ return messages, 0 # -- Optional: per-turn context selection (distinct from compression) -- def select_context( self, request_messages: List[Dict[str, Any]], *, conversation_messages: List[Dict[str, Any]] = None, incoming_message: Dict[str, Any] = None, budget_tokens: int = 0, ) -> List[Dict[str, Any]]: """Optionally *select* (replace) the context for THIS request, pre-generation. Runs every provider request (so also on retries), independent of ``should_compress()``: ``compress()`` shrinks context that is too long, ``select_context()`` swaps in a different context (retrieval, topic routing, branch switching) without abusing ``compress()`` as a per-turn callback. Return ``None`` to leave the request unchanged. The returned list is request-only — it MUST NOT be treated as persisted transcript state; the session DB history is untouched. Unlike the ``pre_llm_call`` hook it may replace the list. The host runs it before prompt cache-control and before every request sanitizer, so a malformed replacement never reaches the provider and the default no-op keeps the request byte-identical (prompt-cache stability preserved). An engine that replaces the list changes its own cache prefix; breakpoints are re-derived on the selected list. ``request_messages`` is the assembled request (system prompt + history + ephemeral prefill); ``conversation_messages`` is the persisted history for reference only (do not mutate); ``budget_tokens`` is the model's context length or 0 if unknown. """ return None def on_turn_complete( self, messages: List[Dict[str, Any]], usage: Dict[str, Any] = None, **kwargs: Any, ) -> None: """Observe a finished turn (complement of ``select_context()``) so the engine can ingest/index/update routing state for the next request. Fires from the normal finalization seam only; some abnormal early returns (content-policy block, provider terminal failure) do not emit it — treat it as best-effort, not guaranteed. ``messages`` is a read-only shallow copy (return value ignored; never rely on transcript mutation). ``usage`` has the ``update_from_response`` dict shape and is ``None`` when the turn never reached a provider response (interrupt). ``kwargs`` may include ``turn_id``, ``task_id``, ``api_call_count``, ``interrupted``, ``failed``, ``turn_exit_reason``. """ return None # -- Optional: pre-flight check ---------------------------------------- def should_compress_preflight(self, messages: List[Dict[str, Any]]) -> bool: """Cheap rough check before the API call (no real token count yet); default skips.""" return False def should_defer_preflight_to_real_usage(self, rough_tokens: int) -> bool: """True when preflight should trust recent real usage over the noisy rough estimate (avoids re-compacting after a compressed request already fit).""" return False def get_automatic_compaction_status_message( self, *, phase: str, default_message: str, **context: Any, ) -> str | None: """User-visible status for automatic compaction, or ``None`` to suppress it. ``phase`` is the host call site (``"preflight"`` / ``"compress"``); ``context`` carries best-effort ``approx_tokens`` / ``threshold_tokens``. Warnings, errors and manual ``/compress`` are not governed by this hook. """ if not self.emit_automatic_compaction_status: return None return default_message # -- Optional: manual /compress preflight ------------------------------ def has_content_to_compress(self, messages: List[Dict[str, Any]]) -> bool: """Preflight guard for gateway ``/compress``: False reports "nothing to compress yet" without an LLM call (e.g. transcript entirely protected).""" return True # -- Optional: session lifecycle --------------------------------------- def on_session_start(self, session_id: str, **kwargs) -> None: """Session begins: load persisted state. kwargs may include hermes_home, platform, model.""" def on_session_end(self, session_id: str, messages: List[Dict[str, Any]]) -> None: """Real session boundary (CLI exit, /reset, gateway expiry) — never per-turn.""" def on_session_reset(self) -> None: """/new or /reset: reset per-session state (default: counters and token tracking).""" self.last_prompt_tokens = 0 self.last_completion_tokens = 0 self.last_total_tokens = 0 self.compression_count = 0 # -- Optional: tools --------------------------------------------------- def get_tool_schemas(self) -> List[Dict[str, Any]]: """Tool schemas this engine exposes to the agent (default: none).""" return [] def handle_tool_call(self, name: str, args: Dict[str, Any], **kwargs) -> str: """Handle a call to one of this engine's tools; must return a JSON string. kwargs may include ``messages`` (live in-memory list).""" import json return json.dumps({"error": f"Unknown context engine tool: {name}"}) # -- Optional: status / display ---------------------------------------- def get_status(self) -> Dict[str, Any]: """Status dict with the standard fields run_agent.py expects.""" # Clamp the -1 "compression just ran, awaiting real usage" sentinel to 0 # so no reader sees a negative usage_percent on the transitional turn. last_prompt = self.last_prompt_tokens if self.last_prompt_tokens > 0 else 0 return { "last_prompt_tokens": last_prompt, "threshold_tokens": self.threshold_tokens, "context_length": self.context_length, "usage_percent": ( min(100, last_prompt / self.context_length * 100) if self.context_length else 0 ), "compression_count": self.compression_count, } # -- Optional: model switch support ------------------------------------ def update_model( self, model: str, context_length: int, base_url: str = "", api_key: str = "", provider: str = "", api_mode: str = "", ) -> None: """Model switch / fallback: recompute threshold_tokens (override for more).""" self.context_length = context_length # Per-model threshold override (longest substring match), else the raw # config percent. Snapshot that percent ONCE so repeated switches fall # back to the configured value, not the previous model's override. from agent.context_compressor import resolve_model_threshold if not hasattr(self, "_config_threshold_percent"): self._config_threshold_percent = self.threshold_percent self._base_threshold_percent = resolve_model_threshold( model, getattr(self, "model_thresholds", {}), self._config_threshold_percent, ) self.threshold_percent = self._base_threshold_percent self.threshold_tokens = int(context_length * self.threshold_percent)