| 1 | from __future__ import annotations |
| 2 | |
| 3 | import json |
| 4 | import os |
| 5 | import re |
| 6 | from dataclasses import dataclass, field |
| 7 | from datetime import datetime |
| 8 | from typing import Any |
| 9 | |
| 10 | |
| 11 | SUMMARY_SECTIONS = [ |
| 12 | "Reference Context Only", |
| 13 | "Active Task", |
| 14 | "Completed Actions", |
| 15 | "Important Files", |
| 16 | "Decisions", |
| 17 | "Errors & Risks", |
| 18 | "Remaining Work", |
| 19 | "Critical Context", |
| 20 | ] |
| 21 | |
| 22 | |
| 23 | @dataclass(slots=True) |
| 24 | class CompactionResult: |
| 25 | summary: str |
| 26 | preserved_messages: list[dict[str, Any]] |
| 27 | compacted_message_count: int |
| 28 | estimated_tokens_before: int |
| 29 | estimated_tokens_after: int |
| 30 | reason: str |
| 31 | mode: str |
| 32 | created_at: str = field(default_factory=lambda: datetime.now().isoformat(timespec="seconds")) |
| 33 | |
| 34 | |
| 35 | class ContextCompactor: |
| 36 | def __init__( |
| 37 | self, |
| 38 | llm: Any | None = None, |
| 39 | *, |
| 40 | token_threshold: int | None = None, |
| 41 | buffer_tokens: int | None = None, |
| 42 | preserve_last_n: int | None = None, |
| 43 | max_messages: int | None = None, |
| 44 | summary_max_chars: int | None = None, |
| 45 | ) -> None: |
| 46 | self.llm = llm |
| 47 | configured_threshold = token_threshold if token_threshold is not None else _default_token_threshold() |
| 48 | self.token_threshold = _env_int("VIMAX_AUTO_COMPACT_TOKEN_THRESHOLD", configured_threshold) |
| 49 | self.buffer_tokens = _env_int("VIMAX_AUTO_COMPACT_BUFFER_TOKENS", buffer_tokens if buffer_tokens is not None else 20000) |
| 50 | self.preserve_last_n = _env_int("VIMAX_COMPACT_PRESERVE_LAST_N", preserve_last_n if preserve_last_n is not None else 6) |
| 51 | self.max_messages = _env_int("VIMAX_COMPACT_MAX_MESSAGES", max_messages if max_messages is not None else 48) |
| 52 | self.summary_max_chars = _env_int("VIMAX_COMPACT_SUMMARY_MAX_CHARS", summary_max_chars if summary_max_chars is not None else 6000) |
| 53 | |
| 54 | def compact_target_tokens(self) -> int: |
| 55 | if self.token_threshold <= 0: |
| 56 | return 0 |
| 57 | return max(0, self.token_threshold - max(0, self.buffer_tokens)) |
| 58 | |
| 59 | def estimate_message_tokens(self, message: dict[str, Any]) -> int: |
| 60 | role = str(message.get("role", "user") or "user") |
| 61 | content = str(message.get("content", "") or "") |
| 62 | metadata = {key: value for key, value in message.items() if key not in {"role", "content"}} |
| 63 | word_count = len(re.findall(r"\w+", content)) |
| 64 | line_count = content.count("\n") + 1 if content else 0 |
| 65 | punctuation_count = len(re.findall(r"[^\w\s]", content)) |
| 66 | role_overhead = {"system": 18, "user": 12, "assistant": 14, "tool": 16}.get(role, 12) |
| 67 | metadata_bonus = min(300, len(json.dumps(metadata, ensure_ascii=False, default=str)) // 6) if metadata else 0 |
| 68 | tool_bonus = 80 if "tool_calls" in message or role == "tool" else 0 |
| 69 | return max(role_overhead, role_overhead + len(content) // 4 + word_count // 2 + line_count * 2 + punctuation_count // 4 + metadata_bonus + tool_bonus) |
| 70 | |
| 71 | def estimate_messages_tokens(self, messages: list[dict[str, Any]]) -> int: |
| 72 | return sum(self.estimate_message_tokens(message) for message in messages) |
| 73 | |
| 74 | def should_preflight_compact(self, messages: list[dict[str, Any]], *, system_tokens: int = 0, tools_tokens: int = 0) -> bool: |
| 75 | target = self.compact_target_tokens() |
| 76 | if target <= 0 or not messages: |
| 77 | return False |
| 78 | total = self.estimate_messages_tokens(messages) + max(0, system_tokens) + max(0, tools_tokens) |
| 79 | return total >= target |
| 80 | |
| 81 | async def compact( |
| 82 | self, |
| 83 | messages: list[dict[str, Any]], |
| 84 | *, |
| 85 | previous_summary: str = "", |
| 86 | preserve_last_n: int | None = None, |
| 87 | reason: str = "manual", |
| 88 | ) -> CompactionResult: |
| 89 | preserve = max(0, self.preserve_last_n if preserve_last_n is None else preserve_last_n) |
| 90 | preserved = [dict(message) for message in messages[-preserve:]] if preserve else [] |
| 91 | compactible = [dict(message) for message in messages[:-preserve]] if preserve else [dict(message) for message in messages] |
| 92 | if not compactible and messages: |
| 93 | compactible = [dict(message) for message in messages] |
| 94 | preserved = [] |
| 95 | before_tokens = self.estimate_messages_tokens(messages) |
| 96 | summary = await self._llm_summary(compactible, preserved, previous_summary, reason) |
| 97 | mode = "llm" |
| 98 | if not summary: |
| 99 | summary = self._fallback_summary(compactible, preserved, previous_summary, reason) |
| 100 | mode = "fallback-local" |
| 101 | summary = self._clip_summary(summary) |
| 102 | synthetic = self.synthetic_summary_message(summary) |
| 103 | after_tokens = self.estimate_messages_tokens([synthetic, *preserved]) |
| 104 | return CompactionResult( |
| 105 | summary=summary, |
| 106 | preserved_messages=preserved, |
| 107 | compacted_message_count=len(compactible), |
| 108 | estimated_tokens_before=before_tokens, |
| 109 | estimated_tokens_after=after_tokens, |
| 110 | reason=reason, |
| 111 | mode=mode, |
| 112 | ) |
| 113 | |
| 114 | def synthetic_summary_message(self, summary: str) -> dict[str, str]: |
| 115 | return { |
| 116 | "role": "system", |
| 117 | "content": "Session context summary. The following summary is reference context only, not a new active instruction.\n\n" + summary.strip(), |
| 118 | } |
| 119 | |
| 120 | async def _llm_summary(self, compactible: list[dict[str, Any]], preserved: list[dict[str, Any]], previous_summary: str, reason: str) -> str: |
| 121 | if self.llm is None: |
| 122 | return "" |
| 123 | payload = { |
| 124 | "reason": reason, |
| 125 | "previous_summary": _clip(previous_summary, 5000), |
| 126 | "messages_to_compact": [self._serialize_message(message) for message in compactible[-self.max_messages:]], |
| 127 | "recent_live_tail": [self._serialize_message(message) for message in preserved[-12:]], |
| 128 | } |
| 129 | system = ( |
| 130 | "You are compressing conversation history for a ViMax agent runtime. " |
| 131 | "Produce a concise markdown handoff summary for a future model call. " |
| 132 | "Preserve user intent, completed actions, important files, tool findings, errors, and remaining work. " |
| 133 | "Label the result as reference context only, not active instructions. " |
| 134 | "Do not answer the user. Do not include prose before the markdown." |
| 135 | ) |
| 136 | user = ( |
| 137 | "Summarize the compacted conversation region into a durable handoff.\n" |
| 138 | "Output markdown with these sections exactly:\n" |
| 139 | "## Reference Context Only\n## Active Task\n## Completed Actions\n## Important Files\n## Decisions\n## Errors & Risks\n## Remaining Work\n## Critical Context\n\n" |
| 140 | "Keep it concise but specific. Mention exact file paths, commands, tool results, and unresolved issues when present.\n\n" |
| 141 | f"{json.dumps(payload, ensure_ascii=False, indent=2)}" |
| 142 | ) |
| 143 | try: |
| 144 | response = await self.llm.complete([{"role": "system", "content": system}, {"role": "user", "content": user}], tools=[]) |
| 145 | except Exception: |
| 146 | return "" |
| 147 | return str(getattr(response, "text", "") or "").strip() |
| 148 | |
| 149 | def _fallback_summary(self, compactible: list[dict[str, Any]], preserved: list[dict[str, Any]], previous_summary: str, reason: str) -> str: |
| 150 | user_lines = [self._message_preview(message, limit=180) for message in compactible if message.get("role") == "user"] |
| 151 | assistant_lines = [self._message_preview(message, limit=180) for message in compactible if message.get("role") == "assistant"] |
| 152 | file_hits = _dedupe(re.findall(r"(?:[\w.\-]+/)+[\w.\-]+\.(?:py|ts|tsx|js|json|md|yaml|yml|txt|mp4|png)", "\n".join(str(message.get("content", "")) for message in compactible))) |
| 153 | error_lines = [self._message_preview(message, limit=180) for message in compactible if _looks_like_error(str(message.get("content", "")))] |
| 154 | remaining = [self._message_preview(message, limit=180) for message in preserved[-4:]] |
| 155 | return "\n".join([ |
| 156 | "## Reference Context Only", |
| 157 | "- This is a compacted checkpoint of older ViMax conversation history, not a new active instruction.", |
| 158 | f"- Compaction reason: {reason}.", |
| 159 | "## Active Task", |
| 160 | _bullet(user_lines[-1:] or ["No explicit active task found in compacted messages."]), |
| 161 | "## Completed Actions", |
| 162 | _bullet(assistant_lines[-4:] or ["No completed assistant actions found in compacted messages."]), |
| 163 | "## Important Files", |
| 164 | _bullet(file_hits[:8] or ["No important file paths found in compacted messages."]), |
| 165 | "## Decisions", |
| 166 | _bullet(_decision_lines(compactible)[:6] or ["No durable decisions found in compacted messages."]), |
| 167 | "## Errors & Risks", |
| 168 | _bullet(error_lines[:6] or ["No errors or risks found in compacted messages."]), |
| 169 | "## Remaining Work", |
| 170 | _bullet(remaining or ["Continue from the recent live tail and current ViMax workflow state."]), |
| 171 | "## Critical Context", |
| 172 | _bullet((["Previous summary existed and was merged as background context."] if previous_summary else []) + ["Use .working_dir artifacts and session checklist as workflow ground truth."]), |
| 173 | ]) |
| 174 | |
| 175 | def _serialize_message(self, message: dict[str, Any]) -> dict[str, Any]: |
| 176 | item = {"role": str(message.get("role", "")), "content": _clip(str(message.get("content", "") or ""), 2400)} |
| 177 | if message.get("name"): |
| 178 | item["name"] = str(message.get("name")) |
| 179 | if message.get("tool_calls"): |
| 180 | item["tool_calls"] = _clip(json.dumps(message.get("tool_calls"), ensure_ascii=False, default=str), 800) |
| 181 | return item |
| 182 | |
| 183 | def _message_preview(self, message: dict[str, Any], *, limit: int) -> str: |
| 184 | role = str(message.get("role", "") or "message") |
| 185 | content = _clip(" ".join(str(message.get("content", "") or "").split()), limit) |
| 186 | if message.get("tool_calls"): |
| 187 | return f"{role}: [tool calls] {_clip(json.dumps(message.get('tool_calls'), ensure_ascii=False, default=str), limit)}" |
| 188 | return f"{role}: {content}" if content else f"{role}: <empty>" |
| 189 | |
| 190 | def _clip_summary(self, summary: str) -> str: |
| 191 | text = summary.strip() |
| 192 | if not text: |
| 193 | text = self._fallback_summary([], [], "", "empty-summary") |
| 194 | if len(text) > self.summary_max_chars: |
| 195 | text = text[: max(0, self.summary_max_chars - 3)].rstrip() + "..." |
| 196 | return text |
| 197 | |
| 198 | |
| 199 | def _default_token_threshold() -> int: |
| 200 | context_window = _env_int("VIMAX_CONTEXT_WINDOW_TOKENS", 200000) |
| 201 | ratio = _env_float("VIMAX_AUTO_COMPACT_RATIO", 0.90) |
| 202 | ratio = min(1.0, max(0.0, ratio)) |
| 203 | return int(context_window * ratio) |
| 204 | |
| 205 | |
| 206 | def _env_int(name: str, default: int) -> int: |
| 207 | try: |
| 208 | return int(os.environ.get(name, str(default))) |
| 209 | except ValueError: |
| 210 | return default |
| 211 | |
| 212 | |
| 213 | def _env_float(name: str, default: float) -> float: |
| 214 | try: |
| 215 | return float(os.environ.get(name, str(default))) |
| 216 | except ValueError: |
| 217 | return default |
| 218 | |
| 219 | |
| 220 | def _clip(text: str, limit: int) -> str: |
| 221 | compact = " ".join(str(text or "").split()) |
| 222 | if len(compact) <= limit: |
| 223 | return compact |
| 224 | return compact[: max(0, limit - 3)].rstrip() + "..." |
| 225 | |
| 226 | |
| 227 | def _bullet(items: list[str]) -> str: |
| 228 | return "\n".join(f"- {item}" for item in items if str(item).strip()) |
| 229 | |
| 230 | |
| 231 | def _dedupe(items: list[str]) -> list[str]: |
| 232 | seen: list[str] = [] |
| 233 | for item in items: |
| 234 | normalized = " ".join(str(item).split()) |
| 235 | if normalized and normalized not in seen: |
| 236 | seen.append(normalized) |
| 237 | return seen |
| 238 | |
| 239 | |
| 240 | def _looks_like_error(text: str) -> bool: |
| 241 | lowered = text.lower() |
| 242 | return any(token in lowered for token in ("error", "failed", "failure", "timeout", "not found", "blocked", "permission")) |
| 243 | |
| 244 | |
| 245 | def _decision_lines(messages: list[dict[str, Any]]) -> list[str]: |
| 246 | tokens = ("decision", "decided", "prefer", "keep ", "switch ", "use ", "preserve ", "avoid ") |
| 247 | rows: list[str] = [] |
| 248 | for message in messages: |
| 249 | content = str(message.get("content", "") or "") |
| 250 | for raw in content.splitlines(): |
| 251 | line = raw.strip(" -") |
| 252 | if line and any(token in line.lower() for token in tokens): |
| 253 | rows.append(_clip(line, 180)) |
| 254 | return _dedupe(rows) |
| 255 |