headroom-ai 0.2.13__py3-none-any.whl

headroom/cache/compression_feedback.py

@@ -0,0 +1,613 @@
+"""Compression Feedback Loop for learning optimal compression strategies.
+
+This module analyzes retrieval patterns from the CompressionStore to learn
+what kinds of compression work well and what doesn't. It provides hints to
+SmartCrusher to improve compression over time.
+
+Key insight from ACON research: Learn compression guidelines by analyzing failures.
+When compression causes the LLM to retrieve more data, that's a signal that
+we compressed too aggressively.
+
+Features:
+- Track retrieval rates per tool type
+- Learn common search queries for each tool
+- Adjust compression aggressiveness based on patterns
+- Provide hints: max_items, fields to preserve, etc.
+
+Usage:
+    feedback = CompressionFeedback(compression_store)
+
+    # Get hints before compressing
+    hints = feedback.get_compression_hints("github_search_repos")
+    # hints = {"max_items": 50, "preserve_fields": ["id", "name"], ...}
+
+    # Apply hints in SmartCrusher config
+    config = SmartCrusherConfig(max_items=hints.get("max_items", 15))
+"""
+
+from __future__ import annotations
+
+import re
+import threading
+import time
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from .compression_store import CompressionStore, RetrievalEvent
+
+
+@dataclass
+class LocalToolPattern:
+    """Learned patterns for a specific tool type (local feedback).
+
+    MEDIUM FIX #18: Renamed from ToolPattern to avoid confusion with
+    headroom.telemetry.toin.ToolPattern which serves a different purpose:
+    - LocalToolPattern: Local feedback patterns keyed by tool_name
+    - toin.ToolPattern: Cross-user TOIN patterns keyed by tool_signature_hash
+    """
+
+    tool_name: str
+
+    # Retrieval statistics
+    total_compressions: int = 0
+    total_retrievals: int = 0
+    full_retrievals: int = 0  # Retrieved entire original content
+    search_retrievals: int = 0  # Used search within content
+
+    # Query analysis
+    common_queries: dict[str, int] = field(default_factory=dict)
+    queried_fields: dict[str, int] = field(default_factory=dict)
+
+    # Strategy analysis - track which strategies work for this tool
+    strategy_compressions: dict[str, int] = field(default_factory=dict)
+    strategy_retrievals: dict[str, int] = field(default_factory=dict)
+
+    # Signature hash tracking - correlate with TOIN patterns
+    signature_hashes: set[str] = field(default_factory=set)
+
+    # Timing
+    last_compression: float = 0.0
+    last_retrieval: float = 0.0
+
+    # Calculated metrics
+    @property
+    def retrieval_rate(self) -> float:
+        """Fraction of compressions that resulted in retrieval."""
+        if self.total_compressions == 0:
+            return 0.0
+        return self.total_retrievals / self.total_compressions
+
+    @property
+    def full_retrieval_rate(self) -> float:
+        """Fraction of retrievals that were full (not search)."""
+        if self.total_retrievals == 0:
+            return 0.0
+        return self.full_retrievals / self.total_retrievals
+
+    @property
+    def search_rate(self) -> float:
+        """Fraction of retrievals that used search."""
+        if self.total_retrievals == 0:
+            return 0.0
+        return self.search_retrievals / self.total_retrievals
+
+    def strategy_retrieval_rate(self, strategy: str) -> float:
+        """Get retrieval rate for a specific compression strategy."""
+        compressions = self.strategy_compressions.get(strategy, 0)
+        if compressions == 0:
+            return 0.0
+        retrievals = self.strategy_retrievals.get(strategy, 0)
+        return retrievals / compressions
+
+    def best_strategy(self) -> str | None:
+        """Find the strategy with lowest retrieval rate (most successful)."""
+        if not self.strategy_compressions:
+            return None
+
+        best = None
+        best_rate = 1.0
+
+        for strategy in self.strategy_compressions:
+            rate = self.strategy_retrieval_rate(strategy)
+            # Only consider strategies with enough samples
+            if self.strategy_compressions[strategy] >= 3 and rate < best_rate:
+                best_rate = rate
+                best = strategy
+
+        return best
+
+
+@dataclass
+class CompressionHints:
+    """Hints for optimizing compression of a specific tool's output."""
+
+    # Item count hints
+    max_items: int = 15  # Default from SmartCrusher
+    min_items: int = 3
+    suggested_items: int | None = None  # Calculated optimal
+
+    # Field preservation
+    preserve_fields: list[str] = field(default_factory=list)
+
+    # Compression aggressiveness (0.0 = aggressive, 1.0 = conservative)
+    aggressiveness: float = 0.7
+
+    # Reasoning
+    reason: str = ""
+
+    # Whether to skip compression entirely
+    skip_compression: bool = False
+
+    # Recommended compression strategy based on local learning
+    recommended_strategy: str | None = None
+
+
+class CompressionFeedback:
+    """Learn from retrieval patterns to improve compression.
+
+    This class analyzes retrieval events from CompressionStore and builds
+    tool-specific patterns. These patterns inform compression decisions.
+
+    Design principles:
+    - High retrieval rate (>50%) → compress less aggressively
+    - Full retrieval dominates → data is unique, skip compression
+    - Search retrieval dominates → keep compressed, add search capability
+    - Frequent queries → preserve fields mentioned in queries
+    """
+
+    # Thresholds for adjusting compression
+    HIGH_RETRIEVAL_THRESHOLD = 0.5  # 50% retrieval = too aggressive
+    MEDIUM_RETRIEVAL_THRESHOLD = 0.2  # 20% retrieval = acceptable
+    MIN_SAMPLES_FOR_HINTS = 5  # Need at least 5 events to make recommendations
+
+    def __init__(
+        self,
+        store: CompressionStore | None = None,
+        enable_learning: bool = True,
+        analysis_interval: float = 60.0,
+    ):
+        """Initialize feedback analyzer.
+
+        Args:
+            store: CompressionStore to analyze. If None, uses global store.
+            enable_learning: Whether to update patterns from events.
+            analysis_interval: Interval in seconds between re-analyzing store events.
+        """
+        self._store = store
+        self._enable_learning = enable_learning
+        self._lock = threading.Lock()
+
+        # Learned patterns per tool
+        self._tool_patterns: dict[str, LocalToolPattern] = {}
+
+        # Time-based tracking
+        self._last_analysis: float = 0.0
+        self._analysis_interval: float = analysis_interval
+        self._last_event_timestamp: float = (
+            0.0  # Track last processed event to avoid double-counting
+        )
+
+        # Global statistics
+        self._total_compressions: int = 0
+        self._total_retrievals: int = 0
+
+    @property
+    def store(self) -> CompressionStore:
+        """Get the compression store (lazy load global if not set)."""
+        if self._store is None:
+            from .compression_store import get_compression_store
+
+            self._store = get_compression_store()
+        return self._store
+
+    def record_compression(
+        self,
+        tool_name: str | None,
+        original_count: int,
+        compressed_count: int,
+        strategy: str | None = None,
+        tool_signature_hash: str | None = None,
+    ) -> None:
+        """Record that a compression occurred.
+
+        Called by SmartCrusher after compressing to track compression events.
+
+        Args:
+            tool_name: Name of the tool whose output was compressed.
+            original_count: Original item count.
+            compressed_count: Compressed item count.
+            strategy: Compression strategy used (e.g., "SMART_SAMPLE", "TOP_N").
+            tool_signature_hash: Hash from ToolSignature for correlation with TOIN.
+        """
+        if not self._enable_learning or not tool_name:
+            return
+
+        with self._lock:
+            self._total_compressions += 1
+
+            if tool_name not in self._tool_patterns:
+                self._tool_patterns[tool_name] = LocalToolPattern(tool_name=tool_name)
+
+            pattern = self._tool_patterns[tool_name]
+            pattern.total_compressions += 1
+            pattern.last_compression = time.time()
+
+            # Track strategy usage
+            if strategy:
+                pattern.strategy_compressions[strategy] = (
+                    pattern.strategy_compressions.get(strategy, 0) + 1
+                )
+
+                # CRITICAL FIX: When truncating strategy dicts, keep them in sync
+                # to prevent desync between compressions and retrievals.
+                # Both dicts must have the same keys for accurate retrieval rate calculation.
+                if len(pattern.strategy_compressions) > 50:
+                    self._truncate_strategy_dicts(pattern)
+
+            # Track signature hash for TOIN correlation
+            if tool_signature_hash:
+                pattern.signature_hashes.add(tool_signature_hash)
+                # CRITICAL FIX: Use deterministic truncation for signature_hashes
+                # Sort lexicographically to ensure consistent behavior across runs
+                if len(pattern.signature_hashes) > 100:
+                    sorted_hashes = sorted(pattern.signature_hashes)[:100]
+                    pattern.signature_hashes = set(sorted_hashes)
+
+    def record_retrieval(
+        self,
+        event: RetrievalEvent,
+        strategy: str | None = None,
+    ) -> None:
+        """Record a retrieval event for pattern learning.
+
+        Called by CompressionStore when content is retrieved.
+
+        Args:
+            event: The retrieval event to record.
+            strategy: Compression strategy that was used (for tracking success rates).
+        """
+        if not self._enable_learning:
+            return
+
+        tool_name = event.tool_name
+        if not tool_name:
+            return
+
+        with self._lock:
+            self._total_retrievals += 1
+
+            if tool_name not in self._tool_patterns:
+                self._tool_patterns[tool_name] = LocalToolPattern(tool_name=tool_name)
+
+            pattern = self._tool_patterns[tool_name]
+            pattern.total_retrievals += 1
+            pattern.last_retrieval = time.time()
+
+            if event.retrieval_type == "full":
+                pattern.full_retrievals += 1
+            else:
+                pattern.search_retrievals += 1
+
+            # Track strategy retrievals (for success rate calculation)
+            if strategy:
+                pattern.strategy_retrievals[strategy] = (
+                    pattern.strategy_retrievals.get(strategy, 0) + 1
+                )
+
+                # CRITICAL FIX: When truncating strategy dicts, keep them in sync
+                # to prevent desync between compressions and retrievals.
+                if len(pattern.strategy_retrievals) > 50:
+                    self._truncate_strategy_dicts(pattern)
+
+            # Track query patterns
+            if event.query:
+                query_lower = event.query.lower()
+                pattern.common_queries[query_lower] = pattern.common_queries.get(query_lower, 0) + 1
+
+                # HIGH: Limit common_queries dict to prevent unbounded growth
+                if len(pattern.common_queries) > 100:
+                    sorted_queries = sorted(
+                        pattern.common_queries.items(),
+                        key=lambda x: x[1],
+                        reverse=True,
+                    )[:100]
+                    pattern.common_queries = dict(sorted_queries)
+
+                # Extract potential field names from query
+                self._extract_field_hints(pattern, event.query)
+
+    def _truncate_strategy_dicts(self, pattern: LocalToolPattern) -> None:
+        """Truncate strategy_compressions and strategy_retrievals in sync.
+
+        CRITICAL FIX: Both dicts must have the same keys for accurate retrieval
+        rate calculation. When truncating, we keep the union of top strategies
+        from both dicts, then truncate both to the same key set.
+        """
+        # Get top 40 strategies from each dict (using 40 to allow union to stay under 50)
+        top_compressions = {
+            k
+            for k, _ in sorted(
+                pattern.strategy_compressions.items(),
+                key=lambda x: x[1],
+                reverse=True,
+            )[:40]
+        }
+        top_retrievals = {
+            k
+            for k, _ in sorted(
+                pattern.strategy_retrievals.items(),
+                key=lambda x: x[1],
+                reverse=True,
+            )[:40]
+        }
+
+        # Keep union of top strategies from both
+        keys_to_keep = top_compressions | top_retrievals
+
+        # Truncate both dicts to same keys
+        pattern.strategy_compressions = {
+            k: v for k, v in pattern.strategy_compressions.items() if k in keys_to_keep
+        }
+        pattern.strategy_retrievals = {
+            k: v for k, v in pattern.strategy_retrievals.items() if k in keys_to_keep
+        }
+
+    def _extract_field_hints(self, pattern: LocalToolPattern, query: str) -> None:
+        """Extract potential field names from search queries.
+
+        Common patterns:
+        - "field:value" or "field=value"
+        - JSON field names like "status", "error", "id"
+        """
+        # Look for field:value patterns
+        field_patterns = re.findall(r"(\w+)[=:]", query)
+        for field_name in field_patterns:
+            pattern.queried_fields[field_name] = pattern.queried_fields.get(field_name, 0) + 1
+
+        # Look for common JSON field names
+        common_fields = [
+            "id",
+            "name",
+            "status",
+            "error",
+            "message",
+            "type",
+            "code",
+            "result",
+            "value",
+            "data",
+            "items",
+            "count",
+        ]
+        query_lower = query.lower()
+        for common_field in common_fields:
+            if common_field in query_lower:
+                pattern.queried_fields[common_field] = (
+                    pattern.queried_fields.get(common_field, 0) + 1
+                )
+
+        # HIGH: Limit queried_fields dict to prevent unbounded growth
+        if len(pattern.queried_fields) > 50:
+            sorted_fields = sorted(
+                pattern.queried_fields.items(),
+                key=lambda x: x[1],
+                reverse=True,
+            )[:50]
+            pattern.queried_fields = dict(sorted_fields)
+
+    def get_compression_hints(
+        self,
+        tool_name: str | None,
+    ) -> CompressionHints:
+        """Get compression hints for a specific tool based on learned patterns.
+
+        Args:
+            tool_name: Name of the tool to get hints for.
+
+        Returns:
+            CompressionHints with recommended settings.
+        """
+        hints = CompressionHints()
+
+        if not tool_name:
+            hints.reason = "No tool name provided, using defaults"
+            return hints
+
+        with self._lock:
+            pattern = self._tool_patterns.get(tool_name)
+
+            if pattern is None:
+                hints.reason = f"No pattern data for {tool_name}, using defaults"
+                return hints
+
+            # Need minimum samples for reliable hints
+            if pattern.total_compressions < self.MIN_SAMPLES_FOR_HINTS:
+                hints.reason = (
+                    f"Insufficient data ({pattern.total_compressions} samples), "
+                    f"need {self.MIN_SAMPLES_FOR_HINTS}"
+                )
+                return hints
+
+            # Calculate hints based on retrieval rate
+            retrieval_rate = pattern.retrieval_rate
+
+            if retrieval_rate > self.HIGH_RETRIEVAL_THRESHOLD:
+                # High retrieval = compress less aggressively
+                if pattern.full_retrieval_rate > 0.8:
+                    # Almost all retrievals are full → skip compression
+                    hints.skip_compression = True
+                    hints.reason = (
+                        f"Very high full retrieval rate ({pattern.full_retrieval_rate:.0%}), "
+                        f"recommending skip compression"
+                    )
+                else:
+                    # Mix of full and search → increase items
+                    hints.max_items = 50
+                    hints.suggested_items = 40
+                    hints.aggressiveness = 0.3
+                    hints.reason = (
+                        f"High retrieval rate ({retrieval_rate:.0%}), "
+                        f"recommending less aggressive compression"
+                    )
+
+            elif retrieval_rate > self.MEDIUM_RETRIEVAL_THRESHOLD:
+                # Medium retrieval = slightly less aggressive
+                hints.max_items = 30
+                hints.suggested_items = 25
+                hints.aggressiveness = 0.5
+                hints.reason = (
+                    f"Medium retrieval rate ({retrieval_rate:.0%}), "
+                    f"recommending moderate compression"
+                )
+
+            else:
+                # Low retrieval = current compression is working
+                hints.max_items = 15
+                hints.suggested_items = 10
+                hints.aggressiveness = 0.7
+                hints.reason = (
+                    f"Low retrieval rate ({retrieval_rate:.0%}), current compression is effective"
+                )
+
+            # Add field preservation hints based on common queries
+            if pattern.queried_fields:
+                # Get top 5 most queried fields
+                sorted_fields = sorted(
+                    pattern.queried_fields.items(),
+                    key=lambda x: x[1],
+                    reverse=True,
+                )[:5]
+                hints.preserve_fields = [f for f, _ in sorted_fields]
+
+            # Recommend the best strategy based on local retrieval patterns
+            best = pattern.best_strategy()
+            if best:
+                hints.recommended_strategy = best
+
+            return hints
+
+    def get_all_patterns(self) -> dict[str, LocalToolPattern]:
+        """Get all learned tool patterns.
+
+        Returns:
+            Dict mapping tool names to their patterns.
+            HIGH FIX: Returns deep copies to prevent external mutation of internal state.
+        """
+        import copy as copy_module
+
+        with self._lock:
+            # Deep copy to prevent external code from modifying internal state
+            return copy_module.deepcopy(self._tool_patterns)
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get feedback statistics for monitoring.
+
+        Returns:
+            Dict with feedback statistics.
+        """
+        with self._lock:
+            return {
+                "total_compressions": self._total_compressions,
+                "total_retrievals": self._total_retrievals,
+                "global_retrieval_rate": (
+                    self._total_retrievals / self._total_compressions
+                    if self._total_compressions > 0
+                    else 0.0
+                ),
+                "tools_tracked": len(self._tool_patterns),
+                "tool_patterns": {
+                    name: {
+                        "compressions": p.total_compressions,
+                        "retrievals": p.total_retrievals,
+                        "retrieval_rate": p.retrieval_rate,
+                        "full_rate": p.full_retrieval_rate,
+                        "search_rate": p.search_rate,
+                        "common_queries": list(p.common_queries.keys())[:5],
+                        "queried_fields": list(p.queried_fields.keys())[:5],
+                    }
+                    for name, p in self._tool_patterns.items()
+                },
+            }
+
+    def analyze_from_store(self) -> None:
+        """Analyze retrieval events from the store.
+
+        This pulls recent events from CompressionStore and updates patterns.
+        Useful for catching up after restart or periodic refresh.
+
+        HIGH FIX: All timestamp reads/writes happen under lock to prevent race
+        conditions where another thread could cause events to be missed or
+        double-counted.
+        """
+        if not self._enable_learning:
+            return
+
+        # Rate limit analysis - check under lock for thread safety
+        now = time.time()
+        with self._lock:
+            if now - self._last_analysis < self._analysis_interval:
+                return
+            # Mark that we're starting analysis (prevents concurrent analysis)
+            self._last_analysis = now
+            last_ts = self._last_event_timestamp
+
+        # Fetch events outside lock (store has its own lock)
+        events = self.store.get_retrieval_events(limit=1000)
+
+        # Filter events to only process new ones (avoid double-counting)
+        new_events = [e for e in events if e.timestamp > last_ts]
+
+        if new_events:
+            # Find the maximum timestamp from new events
+            max_timestamp = max(e.timestamp for e in new_events)
+
+            for event in new_events:
+                self.record_retrieval(event)
+
+            # Update the timestamp AFTER processing - under lock for atomicity
+            with self._lock:
+                # Only update if our max_timestamp is greater than current
+                # (another thread may have processed newer events)
+                if max_timestamp > self._last_event_timestamp:
+                    self._last_event_timestamp = max_timestamp
+
+    def clear(self) -> None:
+        """Clear all learned patterns. Mainly for testing."""
+        with self._lock:
+            self._tool_patterns.clear()
+            self._total_compressions = 0
+            self._total_retrievals = 0
+            self._last_analysis = 0.0
+            self._last_event_timestamp = 0.0
+
+
+# Global feedback instance (lazy initialization)
+_compression_feedback: CompressionFeedback | None = None
+_feedback_lock = threading.Lock()
+
+
+def get_compression_feedback() -> CompressionFeedback:
+    """Get the global compression feedback instance.
+
+    Returns:
+        Global CompressionFeedback instance.
+    """
+    global _compression_feedback
+
+    if _compression_feedback is None:
+        with _feedback_lock:
+            if _compression_feedback is None:
+                _compression_feedback = CompressionFeedback()
+
+    return _compression_feedback
+
+
+def reset_compression_feedback() -> None:
+    """Reset the global compression feedback. Mainly for testing."""
+    global _compression_feedback
+
+    with _feedback_lock:
+        if _compression_feedback is not None:
+            _compression_feedback.clear()
+        _compression_feedback = None