333 lines
16 KiB
Diff
333 lines
16 KiB
Diff
|
From: Timmy Agent <timmy@uniwizard.local>
|
|||
|
|
Date: Mon, 30 Mar 2026 12:43:00 -0700
|
|||
|
|
Subject: [PATCH] Context compression improvements: checkpoints, fallback awareness, validation
|
|||
|
|
|
|||
|
|
This patch addresses critical gaps in the context compressor:
|
|||
|
|
1. Pre-compression checkpoints for recovery
|
|||
|
|
2. Progressive context pressure warnings
|
|||
|
|
3. Summary validation to detect information loss
|
|||
|
|
4. Better tool pruning placeholders
|
|||
|
|
|
|||
|
|
---
|
|||
|
|
agent/context_compressor.py | 102 +++++++++++++++++++++++++++++++++++-
|
|||
|
|
run_agent.py | 71 +++++++++++++++++++++++---
|
|||
|
|
2 files changed, 165 insertions(+), 8 deletions(-)
|
|||
|
|
|
|||
|
|
diff --git a/agent/context_compressor.py b/agent/context_compressor.py
|
|||
|
|
index abc123..def456 100644
|
|||
|
|
--- a/agent/context_compressor.py
|
|||
|
|
+++ b/agent/context_compressor.py
|
|||
|
|
@@ -15,6 +15,7 @@ Improvements over v1:
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
from typing import Any, Dict, List, Optional
|
|||
|
|
+import re
|
|||
|
|
|
|||
|
|
from agent.auxiliary_client import call_llm
|
|||
|
|
from agent.model_metadata import (
|
|||
|
|
@@ -44,6 +45,12 @@ _SUMMARY_TOKENS_CEILING = 8000
|
|||
|
|
# Placeholder used when pruning old tool results
|
|||
|
|
_PRUNED_TOOL_PLACEHOLDER = "[Old tool output cleared to save context space]"
|
|||
|
|
|
|||
|
|
+# Enhanced placeholder with context (used when we know the tool name)
|
|||
|
|
+_PRUNED_TOOL_PLACEHOLDER_TEMPLATE = (
|
|||
|
|
+ "[Tool output for '{tool_name}' cleared to save context space. "
|
|||
|
|
+ "Original: {original_chars} chars]"
|
|||
|
|
+)
|
|||
|
|
+
|
|||
|
|
# Chars per token rough estimate
|
|||
|
|
_CHARS_PER_TOKEN = 4
|
|||
|
|
|
|||
|
|
@@ -152,13 +159,22 @@ class ContextCompressor:
|
|||
|
|
def _prune_old_tool_results(
|
|||
|
|
self, messages: List[Dict[str, Any]], protect_tail_count: int,
|
|||
|
|
) -> tuple[List[Dict[str, Any]], int]:
|
|||
|
|
- """Replace old tool result contents with a short placeholder.
|
|||
|
|
+ """Replace old tool result contents with an informative placeholder.
|
|||
|
|
|
|||
|
|
Walks backward from the end, protecting the most recent
|
|||
|
|
- ``protect_tail_count`` messages. Older tool results get their
|
|||
|
|
- content replaced with a placeholder string.
|
|||
|
|
+ ``protect_tail_count`` messages. Older tool results are summarized
|
|||
|
|
+ with an informative placeholder that includes the tool name.
|
|||
|
|
|
|||
|
|
Returns (pruned_messages, pruned_count).
|
|||
|
|
+
|
|||
|
|
+ Improvement: Now includes tool name in placeholder for better
|
|||
|
|
+ context about what was removed.
|
|||
|
|
"""
|
|||
|
|
if not messages:
|
|||
|
|
return messages, 0
|
|||
|
|
@@ -170,10 +186,26 @@ class ContextCompressor:
|
|||
|
|
for i in range(prune_boundary):
|
|||
|
|
msg = result[i]
|
|||
|
|
if msg.get("role") != "tool":
|
|||
|
|
continue
|
|||
|
|
content = msg.get("content", "")
|
|||
|
|
if not content or content == _PRUNED_TOOL_PLACEHOLDER:
|
|||
|
|
continue
|
|||
|
|
# Only prune if the content is substantial (>200 chars)
|
|||
|
|
if len(content) > 200:
|
|||
|
|
- result[i] = {**msg, "content": _PRUNED_TOOL_PLACEHOLDER}
|
|||
|
|
+ # Try to find the tool name from the matching assistant message
|
|||
|
|
+ tool_call_id = msg.get("tool_call_id", "")
|
|||
|
|
+ tool_name = "unknown"
|
|||
|
|
+ for m in messages:
|
|||
|
|
+ if m.get("role") == "assistant" and m.get("tool_calls"):
|
|||
|
|
+ for tc in m.get("tool_calls", []):
|
|||
|
|
+ tc_id = tc.get("id", "") if isinstance(tc, dict) else getattr(tc, "id", "")
|
|||
|
|
+ if tc_id == tool_call_id:
|
|||
|
|
+ fn = tc.get("function", {}) if isinstance(tc, dict) else getattr(tc, "function", {})
|
|||
|
|
+ tool_name = fn.get("name", "unknown") if isinstance(fn, dict) else getattr(fn, "name", "unknown")
|
|||
|
|
+ break
|
|||
|
|
+
|
|||
|
|
+ placeholder = _PRUNED_TOOL_PLACEHOLDER_TEMPLATE.format(
|
|||
|
|
+ tool_name=tool_name,
|
|||
|
|
+ original_chars=len(content)
|
|||
|
|
+ )
|
|||
|
|
+ result[i] = {**msg, "content": placeholder}
|
|||
|
|
pruned += 1
|
|||
|
|
|
|||
|
|
return result, pruned
|
|||
|
|
@@ -250,6 +282,52 @@ class ContextCompressor:
|
|||
|
|
## Critical Context
|
|||
|
|
[Any specific values, error messages, configuration details, or data that would be lost without explicit preservation]
|
|||
|
|
|
|||
|
|
+ Target ~{summary_budget} tokens. Be specific — include file paths, command outputs, error messages, and concrete values rather than vague descriptions.
|
|||
|
|
+
|
|||
|
|
+ Write only the summary body. Do not include any preamble or prefix."""
|
|||
|
|
+
|
|||
|
|
+ def _extract_critical_refs(self, turns: List[Dict[str, Any]]) -> set:
|
|||
|
|
+ """Extract critical references that should appear in a valid summary.
|
|||
|
|
+
|
|||
|
|
+ Returns set of file paths, error signatures, and key values that
|
|||
|
|
+ the summary should preserve.
|
|||
|
|
+ """
|
|||
|
|
+ critical = set()
|
|||
|
|
+ for msg in turns:
|
|||
|
|
+ content = msg.get("content", "") or ""
|
|||
|
|
+ if not isinstance(content, str):
|
|||
|
|
+ continue
|
|||
|
|
+
|
|||
|
|
+ # File paths (common code extensions)
|
|||
|
|
+ for match in re.finditer(r'[\w\-./]+\.(py|js|ts|jsx|tsx|json|yaml|yml|md|txt|rs|go|java|cpp|c|h|hpp)\b', content):
|
|||
|
|
+ critical.add(match.group(0))
|
|||
|
|
+
|
|||
|
|
+ # Error patterns
|
|||
|
|
+ lines = content.split('\n')
|
|||
|
|
+ for line in lines:
|
|||
|
|
+ line_lower = line.lower()
|
|||
|
|
+ if any(k in line_lower for k in ['error:', 'exception:', 'traceback', 'failed:', 'failure:']):
|
|||
|
|
+ # First 80 chars of error line
|
|||
|
|
+ critical.add(line[:80].strip())
|
|||
|
|
+
|
|||
|
|
+ # URLs
|
|||
|
|
+ for match in re.finditer(r'https?://[^\s<>"\']+', content):
|
|||
|
|
+ critical.add(match.group(0))
|
|||
|
|
+
|
|||
|
|
+ return critical
|
|||
|
|
+
|
|||
|
|
+ def _validate_summary(self, summary: str, turns: List[Dict[str, Any]]) -> tuple[bool, List[str]]:
|
|||
|
|
+ """Validate that summary captures critical information from turns.
|
|||
|
|
+
|
|||
|
|
+ Returns (is_valid, missing_critical_items).
|
|||
|
|
+ """
|
|||
|
|
+ if not summary or len(summary) < 50:
|
|||
|
|
+ return False, ["summary too short"]
|
|||
|
|
+
|
|||
|
|
+ critical = self._extract_critical_refs(turns)
|
|||
|
|
+ if not critical:
|
|||
|
|
+ return True, []
|
|||
|
|
+
|
|||
|
|
+ # Check what critical items are missing from summary
|
|||
|
|
+ missing = [ref for ref in critical if ref not in summary]
|
|||
|
|
+
|
|||
|
|
+ # Allow up to 50% loss of non-critical references
|
|||
|
|
+ if len(missing) > len(critical) * 0.5 and len(critical) > 3:
|
|||
|
|
+ return False, missing[:5] # Return first 5 missing items
|
|||
|
|
+
|
|||
|
|
+ return True, []
|
|||
|
|
+
|
|||
|
|
+ def _generate_summary(self, turns_to_summarize: List[Dict[str, Any]]) -> Optional[str]:
|
|||
|
|
+ """Generate a structured summary of conversation turns.
|
|||
|
|
+
|
|||
|
|
+ NEW: Added validation step to detect low-quality summaries.
|
|||
|
|
+ Falls back to extended summarization if validation fails.
|
|||
|
|
+ """
|
|||
|
|
+ summary_budget = self._compute_summary_budget(turns_to_summarize)
|
|||
|
|
+ content_to_summarize = self._serialize_for_summary(turns_to_summarize)
|
|||
|
|
+
|
|||
|
|
if self._previous_summary:
|
|||
|
|
# Iterative update: preserve existing info, add new progress
|
|||
|
|
prompt = f"""You are updating a context compaction summary...
|
|||
|
|
@@ -341,9 +419,27 @@ class ContextCompressor:
|
|||
|
|
try:
|
|||
|
|
call_kwargs = {
|
|||
|
|
"task": "compression",
|
|||
|
|
"messages": [{"role": "user", "content": prompt}],
|
|||
|
|
"temperature": 0.3,
|
|||
|
|
- "max_tokens": summary_budget * 2,
|
|||
|
|
+ "max_tokens": min(summary_budget * 2, 4000),
|
|||
|
|
}
|
|||
|
|
if self.summary_model:
|
|||
|
|
call_kwargs["model"] = self.summary_model
|
|||
|
|
response = call_llm(**call_kwargs)
|
|||
|
|
content = response.choices[0].message.content
|
|||
|
|
# Handle cases where content is not a string (e.g., dict from llama.cpp)
|
|||
|
|
if not isinstance(content, str):
|
|||
|
|
content = str(content) if content else ""
|
|||
|
|
summary = content.strip()
|
|||
|
|
+
|
|||
|
|
+ # NEW: Validate the generated summary
|
|||
|
|
+ is_valid, missing = self._validate_summary(summary, turns_to_summarize)
|
|||
|
|
+ if not is_valid and not self.quiet_mode:
|
|||
|
|
+ logger.warning(
|
|||
|
|
+ "Summary validation detected potential information loss. "
|
|||
|
|
+ "Missing: %s", missing
|
|||
|
|
+ )
|
|||
|
|
+ # Attempt to extend the summary with missing critical info
|
|||
|
|
+ if missing:
|
|||
|
|
+ critical_note = "\n\n## Critical Items Preserved\n" + "\n".join(f"- {m}" for m in missing[:10])
|
|||
|
|
+ summary = summary + critical_note
|
|||
|
|
+
|
|||
|
|
# Store for iterative updates on next compaction
|
|||
|
|
self._previous_summary = summary
|
|||
|
|
return self._with_summary_prefix(summary)
|
|||
|
|
@@ -660,6 +756,10 @@ class ContextCompressor:
|
|||
|
|
saved_estimate,
|
|||
|
|
)
|
|||
|
|
logger.info("Compression #%d complete", self.compression_count)
|
|||
|
|
+
|
|||
|
|
+ # NEW: Log compression efficiency metric
|
|||
|
|
+ if display_tokens > 0:
|
|||
|
|
+ efficiency = saved_estimate / display_tokens * 100
|
|||
|
|
+ logger.info("Compression efficiency: %.1f%% tokens saved", efficiency)
|
|||
|
|
|
|||
|
|
return compressed
|
|||
|
|
|
|||
|
|
diff --git a/run_agent.py b/run_agent.py
|
|||
|
|
index abc123..def456 100644
|
|||
|
|
--- a/run_agent.py
|
|||
|
|
+++ b/run_agent.py
|
|||
|
|
@@ -1186,7 +1186,35 @@ class AIAgent:
|
|||
|
|
pass
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
+ # NEW: Collect context lengths for all models in fallback chain
|
|||
|
|
+ # This ensures compression threshold is appropriate for ANY model that might be used
|
|||
|
|
+ _fallback_contexts = []
|
|||
|
|
+ _fallback_providers = _agent_cfg.get("fallback_providers", [])
|
|||
|
|
+ if isinstance(_fallback_providers, list):
|
|||
|
|
+ for fb in _fallback_providers:
|
|||
|
|
+ if isinstance(fb, dict):
|
|||
|
|
+ fb_model = fb.get("model", "")
|
|||
|
|
+ fb_base = fb.get("base_url", "")
|
|||
|
|
+ fb_provider = fb.get("provider", "")
|
|||
|
|
+ fb_key_env = fb.get("api_key_env", "")
|
|||
|
|
+ fb_key = os.getenv(fb_key_env, "")
|
|||
|
|
+ if fb_model:
|
|||
|
|
+ try:
|
|||
|
|
+ fb_ctx = get_model_context_length(
|
|||
|
|
+ fb_model, base_url=fb_base,
|
|||
|
|
+ api_key=fb_key, provider=fb_provider
|
|||
|
|
+ )
|
|||
|
|
+ if fb_ctx and fb_ctx > 0:
|
|||
|
|
+ _fallback_contexts.append(fb_ctx)
|
|||
|
|
+ except Exception:
|
|||
|
|
+ pass
|
|||
|
|
+
|
|||
|
|
+ # Use minimum context length for conservative compression
|
|||
|
|
+ # This ensures we compress early enough for the most constrained model
|
|||
|
|
+ _effective_context_length = _config_context_length
|
|||
|
|
+ if _fallback_contexts:
|
|||
|
|
+ _min_fallback = min(_fallback_contexts)
|
|||
|
|
+ if _effective_context_length is None or _min_fallback < _effective_context_length:
|
|||
|
|
+ _effective_context_length = _min_fallback
|
|||
|
|
+ if not self.quiet_mode:
|
|||
|
|
+ print(f"📊 Using conservative context limit: {_effective_context_length:,} tokens (fallback-aware)")
|
|||
|
|
+
|
|||
|
|
self.context_compressor = ContextCompressor(
|
|||
|
|
model=self.model,
|
|||
|
|
threshold_percent=compression_threshold,
|
|||
|
|
@@ -1196,7 +1224,7 @@ class AIAgent:
|
|||
|
|
summary_model_override=compression_summary_model,
|
|||
|
|
quiet_mode=self.quiet_mode,
|
|||
|
|
base_url=self.base_url,
|
|||
|
|
api_key=self.api_key,
|
|||
|
|
- config_context_length=_config_context_length,
|
|||
|
|
+ config_context_length=_effective_context_length,
|
|||
|
|
provider=self.provider,
|
|||
|
|
)
|
|||
|
|
self.compression_enabled = compression_enabled
|
|||
|
|
@@ -5248,6 +5276,22 @@ class AIAgent:
|
|||
|
|
|
|||
|
|
def _compress_context(self, messages: list, system_message: str, *, approx_tokens: int = None, task_id: str = "default") -> tuple:
|
|||
|
|
"""Compress conversation context and split the session in SQLite.
|
|||
|
|
+
|
|||
|
|
+ NEW: Creates a checkpoint before compression for recovery.
|
|||
|
|
+ This allows rewinding if the summary loses critical information.
|
|||
|
|
+
|
|||
|
|
+ Checkpoint naming: pre-compression-N where N is compression count
|
|||
|
|
+ The checkpoint is kept for potential recovery but marked as internal.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
(compressed_messages, new_system_prompt) tuple
|
|||
|
|
"""
|
|||
|
|
+ # NEW: Create checkpoint BEFORE compression
|
|||
|
|
+ if self._checkpoint_mgr and hasattr(self._checkpoint_mgr, 'create_checkpoint'):
|
|||
|
|
+ try:
|
|||
|
|
+ checkpoint_name = f"pre-compression-{self.context_compressor.compression_count}"
|
|||
|
|
+ self._checkpoint_mgr.create_checkpoint(
|
|||
|
|
+ name=checkpoint_name,
|
|||
|
|
+ description=f"Automatic checkpoint before compression #{self.context_compressor.compression_count}"
|
|||
|
|
+ )
|
|||
|
|
+ if not self.quiet_mode:
|
|||
|
|
+ logger.info(f"Created checkpoint '{checkpoint_name}' before compression")
|
|||
|
|
+ except Exception as e:
|
|||
|
|
+ logger.debug(f"Failed to create pre-compression checkpoint: {e}")
|
|||
|
|
+
|
|||
|
|
# Pre-compression memory flush: let the model save memories before they're lost
|
|||
|
|
self.flush_memories(messages, min_turns=0)
|
|||
|
|
|
|||
|
|
@@ -7862,12 +7906,33 @@ class AIAgent:
|
|||
|
|
# Update compressor with actual token count for accurate threshold check
|
|||
|
|
if hasattr(self, 'context_compressor') and self.context_compressor:
|
|||
|
|
self.context_compressor.update_from_response(usage_dict)
|
|||
|
|
- # Show context pressure warning at 85% of compaction threshold
|
|||
|
|
+
|
|||
|
|
+ # NEW: Progressive context pressure warnings
|
|||
|
|
_compressor = self.context_compressor
|
|||
|
|
if _compressor.threshold_tokens > 0:
|
|||
|
|
_compaction_progress = _real_tokens / _compressor.threshold_tokens
|
|||
|
|
- if _compaction_progress >= 0.85 and not self._context_pressure_warned:
|
|||
|
|
- self._context_pressure_warned = True
|
|||
|
|
+
|
|||
|
|
+ # Progressive warning levels
|
|||
|
|
+ _warning_levels = [
|
|||
|
|
+ (0.60, "info", "ℹ️ Context usage at 60%"),
|
|||
|
|
+ (0.75, "notice", "📊 Context usage at 75% — consider wrapping up"),
|
|||
|
|
+ (0.85, "warning", "⚠️ Context usage at 85% — compression imminent"),
|
|||
|
|
+ (0.95, "critical", "🔴 Context usage at 95% — compression will trigger soon"),
|
|||
|
|
+ ]
|
|||
|
|
+
|
|||
|
|
+ if not hasattr(self, '_context_pressure_reported'):
|
|||
|
|
+ self._context_pressure_reported = set()
|
|||
|
|
+
|
|||
|
|
+ for threshold, level, message in _warning_levels:
|
|||
|
|
+ if _compaction_progress >= threshold and threshold not in self._context_pressure_reported:
|
|||
|
|
+ self._context_pressure_reported.add(threshold)
|
|||
|
|
+ # Only show warnings at 85%+ in quiet mode
|
|||
|
|
+ if level in ("warning", "critical") or not self.quiet_mode:
|
|||
|
|
+ if self.status_callback:
|
|||
|
|
+ self.status_callback(level, message)
|
|||
|
|
+ print(f"\n{message}\n")
|
|||
|
|
+
|
|||
|
|
+ # Legacy single warning for backward compatibility
|
|||
|
|
+ if _compaction_progress >= 0.85 and not getattr(self, '_context_pressure_warned', False):
|
|||
|
|
+ self._context_pressure_warned = True # Mark legacy flag
|
|||
|
|
_ctx_msg = (
|
|||
|
|
f"📊 Context is at {_compaction_progress:.0%} of compression threshold "
|
|||
|
|
f"({_real_tokens:,} / {_compressor.threshold_tokens:,} tokens). "
|
|||
|
|
--
|
|||
|
|
2.40.0
|