| """ContextManager — token-budgeted compaction window (ml-intern pattern). |
| |
| Tracks per-message token estimates, summarizes oldest turns when the window |
| exceeds `target_tokens`, preserves the system prompt and the most recent |
| user/tool turns verbatim. Pure stdlib; optional summarizer callback. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Callable |
| from typing import Any |
|
|
| from pydantic import BaseModel, ConfigDict, Field |
|
|
|
|
| class ContextManagerConfig(BaseModel): |
| model_config = ConfigDict(extra="forbid") |
|
|
| target_tokens: int = Field(default=170_000, ge=1024) |
| keep_recent: int = Field( |
| default=8, |
| ge=1, |
| description="messages from the end to keep verbatim", |
| ) |
| chars_per_token: float = Field(default=4.0, gt=0.0) |
|
|
|
|
| def estimate_tokens(text: str, chars_per_token: float = 4.0) -> int: |
| """Cheap whitespace/char-rate token estimate (no tokenizer dep).""" |
| if not text: |
| return 0 |
| |
| return max(1, int(len(text) / chars_per_token)) |
|
|
|
|
| def message_tokens(message: dict[str, Any], chars_per_token: float = 4.0) -> int: |
| """Estimate tokens for one chat message.""" |
| text_parts: list[str] = [str(message.get("role", "")), str(message.get("content", "") or "")] |
| for c in message.get("tool_calls") or []: |
| fn = c.get("function", {}) |
| text_parts.append(str(fn.get("name", ""))) |
| text_parts.append(str(fn.get("arguments", ""))) |
| return estimate_tokens(" ".join(text_parts), chars_per_token) |
|
|
|
|
| SummarizeFn = Callable[[list[dict[str, Any]]], str] |
|
|
|
|
| class ContextManager: |
| """Trim message lists to fit `target_tokens`.""" |
|
|
| def __init__( |
| self, |
| cfg: ContextManagerConfig | None = None, |
| summarize: SummarizeFn | None = None, |
| ) -> None: |
| self.cfg = cfg or ContextManagerConfig() |
| self.summarize = summarize |
|
|
| def estimate(self, messages: list[dict[str, Any]]) -> int: |
| return sum(message_tokens(m, self.cfg.chars_per_token) for m in messages) |
|
|
| def compact(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: |
| """Return a trimmed message list whose token estimate is <= target_tokens. |
| |
| Strategy: |
| 1. Always keep the leading system message (if present) verbatim. |
| 2. Always keep the trailing `keep_recent` messages verbatim. |
| 3. The middle band is replaced with a single summary turn (via |
| `summarize` callback) or simply elided with a meta note. |
| """ |
| if not messages: |
| return [] |
|
|
| if self.estimate(messages) <= self.cfg.target_tokens: |
| return list(messages) |
|
|
| head: list[dict[str, Any]] = [] |
| if messages and messages[0].get("role") == "system": |
| head = [messages[0]] |
| body = messages[1:] |
| else: |
| body = list(messages) |
|
|
| if len(body) <= self.cfg.keep_recent: |
| return head + body |
|
|
| recent = body[-self.cfg.keep_recent :] |
| elided = body[: -self.cfg.keep_recent] |
|
|
| if self.summarize is not None: |
| summary_text = self.summarize(elided) |
| else: |
| summary_text = ( |
| f"[context compacted: {len(elided)} earlier messages elided" |
| f" to fit {self.cfg.target_tokens}-token window]" |
| ) |
|
|
| summary_msg = {"role": "system", "content": summary_text} |
| compacted = [*head, summary_msg, *recent] |
|
|
| |
| if self.estimate(compacted) > self.cfg.target_tokens and self.cfg.keep_recent > 1: |
| tighter_cfg = self.cfg.model_copy(update={"keep_recent": max(1, self.cfg.keep_recent // 2)}) |
| return ContextManager(tighter_cfg, self.summarize).compact(messages) |
|
|
| return compacted |
|
|
|
|
| __all__ = ["ContextManager", "ContextManagerConfig", "estimate_tokens", "message_tokens"] |
|
|
