alma-memory 0.5.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

alma/retrieval/budget.py

@@ -0,0 +1,525 @@
+"""
+ALMA Token Budget Management.
+
+Implements attention budget tracking for retrieval to prevent context overflow.
+Based on context engineering principles: "Context window space is finite and expensive."
+
+Features:
+- Token counting with configurable estimator
+- Priority-based inclusion (must-see → should-see → fetch-on-demand)
+- Budget enforcement with graceful degradation
+- Tracking and metrics for optimization
+"""
+
+import logging
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+from alma.types import (
+    AntiPattern,
+    DomainKnowledge,
+    Heuristic,
+    MemorySlice,
+    Outcome,
+    UserPreference,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class PriorityTier(Enum):
+    """
+    Priority tiers for attention budget allocation.
+
+    Based on context engineering guidance:
+    - MUST_SEE: In the window for every step, no exceptions
+    - SHOULD_SEE: Important but can be summarized if needed
+    - FETCH_ON_DEMAND: Referenced but not included; agent can request
+    - EXCLUDE: Never read again; processed and compressed
+    """
+
+    MUST_SEE = 1
+    SHOULD_SEE = 2
+    FETCH_ON_DEMAND = 3
+    EXCLUDE = 4
+
+
+@dataclass
+class BudgetConfig:
+    """Configuration for token budget management."""
+
+    # Total budget in tokens
+    max_tokens: int = 4000
+
+    # Per-tier allocation (percentages of max_tokens)
+    must_see_pct: float = 0.4  # 40% for critical memories
+    should_see_pct: float = 0.35  # 35% for important memories
+    fetch_on_demand_pct: float = 0.25  # 25% reserved for on-demand
+
+    # Per-type limits (within tier allocation)
+    max_heuristics: int = 10
+    max_outcomes: int = 10
+    max_knowledge: int = 5
+    max_anti_patterns: int = 5
+    max_preferences: int = 5
+
+    # Token estimation
+    chars_per_token: int = 4  # Rough estimate: 4 chars = 1 token
+
+    # Truncation settings
+    truncate_long_content: bool = True
+    max_content_chars: int = 500  # Truncate individual items
+
+    def get_tier_budget(self, tier: PriorityTier) -> int:
+        """Get token budget for a priority tier."""
+        if tier == PriorityTier.MUST_SEE:
+            return int(self.max_tokens * self.must_see_pct)
+        elif tier == PriorityTier.SHOULD_SEE:
+            return int(self.max_tokens * self.should_see_pct)
+        elif tier == PriorityTier.FETCH_ON_DEMAND:
+            return int(self.max_tokens * self.fetch_on_demand_pct)
+        return 0
+
+
+@dataclass
+class BudgetedItem:
+    """A memory item with budget metadata."""
+
+    item: Any
+    memory_type: str
+    priority: PriorityTier
+    estimated_tokens: int
+    included: bool = True
+    truncated: bool = False
+    summary_only: bool = False
+
+
+@dataclass
+class BudgetReport:
+    """Report on budget usage after retrieval."""
+
+    total_budget: int
+    used_tokens: int
+    remaining_tokens: int
+
+    # Per-tier breakdown
+    must_see_used: int = 0
+    must_see_budget: int = 0
+    should_see_used: int = 0
+    should_see_budget: int = 0
+    fetch_on_demand_count: int = 0
+
+    # Items by status
+    included_count: int = 0
+    excluded_count: int = 0
+    truncated_count: int = 0
+    summary_only_count: int = 0
+
+    # Warnings
+    budget_exceeded: bool = False
+    items_dropped: List[str] = field(default_factory=list)
+
+    @property
+    def utilization_pct(self) -> float:
+        """Budget utilization percentage."""
+        if self.total_budget == 0:
+            return 0.0
+        return (self.used_tokens / self.total_budget) * 100
+
+
+class TokenEstimator:
+    """Estimates token count for memory items."""
+
+    def __init__(self, chars_per_token: int = 4):
+        self.chars_per_token = chars_per_token
+
+    def estimate(self, item: Any) -> int:
+        """Estimate tokens for any memory item."""
+        if isinstance(item, Heuristic):
+            return self._estimate_heuristic(item)
+        elif isinstance(item, Outcome):
+            return self._estimate_outcome(item)
+        elif isinstance(item, DomainKnowledge):
+            return self._estimate_knowledge(item)
+        elif isinstance(item, AntiPattern):
+            return self._estimate_anti_pattern(item)
+        elif isinstance(item, UserPreference):
+            return self._estimate_preference(item)
+        elif isinstance(item, str):
+            return len(item) // self.chars_per_token
+        return 50  # Default estimate
+
+    def _estimate_heuristic(self, h: Heuristic) -> int:
+        """Estimate tokens for a heuristic."""
+        text = f"{h.condition} {h.strategy}"
+        return len(text) // self.chars_per_token + 20  # +20 for metadata
+
+    def _estimate_outcome(self, o: Outcome) -> int:
+        """Estimate tokens for an outcome."""
+        text = f"{o.task_type} {o.task_description} {o.strategy_used}"
+        if o.error_message:
+            text += f" {o.error_message}"
+        return len(text) // self.chars_per_token + 15
+
+    def _estimate_knowledge(self, k: DomainKnowledge) -> int:
+        """Estimate tokens for domain knowledge."""
+        # Knowledge can have complex 'fact' structures
+        fact_str = str(k.fact) if k.fact else ""
+        text = f"{k.domain} {fact_str}"
+        return len(text) // self.chars_per_token + 10
+
+    def _estimate_anti_pattern(self, ap: AntiPattern) -> int:
+        """Estimate tokens for an anti-pattern."""
+        text = f"{ap.pattern} {ap.why_bad} {ap.better_alternative}"
+        return len(text) // self.chars_per_token + 15
+
+    def _estimate_preference(self, p: UserPreference) -> int:
+        """Estimate tokens for a preference."""
+        text = f"{p.category} {p.preference} {p.context or ''}"
+        return len(text) // self.chars_per_token + 10
+
+    def estimate_slice(self, memory_slice: MemorySlice) -> int:
+        """Estimate total tokens for a MemorySlice."""
+        total = 0
+        for h in memory_slice.heuristics:
+            total += self._estimate_heuristic(h)
+        for o in memory_slice.outcomes:
+            total += self._estimate_outcome(o)
+        for k in memory_slice.domain_knowledge:
+            total += self._estimate_knowledge(k)
+        for ap in memory_slice.anti_patterns:
+            total += self._estimate_anti_pattern(ap)
+        for p in memory_slice.preferences:
+            total += self._estimate_preference(p)
+        return total
+
+
+class RetrievalBudget:
+    """
+    Manages token budget for memory retrieval.
+
+    Ensures retrieval results fit within context window limits
+    while prioritizing the most important memories.
+
+    Usage:
+        budget = RetrievalBudget(config=BudgetConfig(max_tokens=4000))
+
+        # Check before including
+        if budget.can_include(memory, priority=PriorityTier.MUST_SEE):
+            budget.include(memory, "heuristic", PriorityTier.MUST_SEE)
+
+        # Or process entire slice
+        budgeted_slice, report = budget.apply_budget(memory_slice, priorities)
+    """
+
+    def __init__(
+        self,
+        config: Optional[BudgetConfig] = None,
+        estimator: Optional[TokenEstimator] = None,
+        priority_classifier: Optional[Callable[[Any, str], PriorityTier]] = None,
+    ):
+        self.config = config or BudgetConfig()
+        self.estimator = estimator or TokenEstimator(self.config.chars_per_token)
+        self.priority_classifier = priority_classifier or self._default_classifier
+
+        # Tracking
+        self._used_tokens = 0
+        self._tier_usage: Dict[PriorityTier, int] = {
+            PriorityTier.MUST_SEE: 0,
+            PriorityTier.SHOULD_SEE: 0,
+            PriorityTier.FETCH_ON_DEMAND: 0,
+        }
+        self._items: List[BudgetedItem] = []
+        self._excluded: List[str] = []
+
+    def reset(self) -> None:
+        """Reset budget tracking for new retrieval."""
+        self._used_tokens = 0
+        self._tier_usage = {
+            PriorityTier.MUST_SEE: 0,
+            PriorityTier.SHOULD_SEE: 0,
+            PriorityTier.FETCH_ON_DEMAND: 0,
+        }
+        self._items = []
+        self._excluded = []
+
+    @property
+    def remaining_tokens(self) -> int:
+        """Tokens remaining in total budget."""
+        return max(0, self.config.max_tokens - self._used_tokens)
+
+    @property
+    def used_tokens(self) -> int:
+        """Tokens used so far."""
+        return self._used_tokens
+
+    def can_include(
+        self,
+        item: Any,
+        priority: PriorityTier = PriorityTier.SHOULD_SEE,
+    ) -> bool:
+        """Check if an item can be included within budget."""
+        if priority == PriorityTier.EXCLUDE:
+            return False
+
+        estimated = self.estimator.estimate(item)
+        tier_budget = self.config.get_tier_budget(priority)
+        tier_used = self._tier_usage.get(priority, 0)
+
+        # Check tier budget
+        if tier_used + estimated > tier_budget:
+            return False
+
+        # Check total budget
+        if self._used_tokens + estimated > self.config.max_tokens:
+            return False
+
+        return True
+
+    def include(
+        self,
+        item: Any,
+        memory_type: str,
+        priority: PriorityTier = PriorityTier.SHOULD_SEE,
+        force: bool = False,
+    ) -> BudgetedItem:
+        """
+        Include an item in the budget.
+
+        Args:
+            item: Memory item to include
+            memory_type: Type name (heuristic, outcome, etc.)
+            priority: Priority tier for allocation
+            force: Include even if over budget (for MUST_SEE items)
+
+        Returns:
+            BudgetedItem with inclusion status
+        """
+        estimated = self.estimator.estimate(item)
+        can_fit = self.can_include(item, priority)
+
+        budgeted = BudgetedItem(
+            item=item,
+            memory_type=memory_type,
+            priority=priority,
+            estimated_tokens=estimated,
+            included=can_fit or force,
+        )
+
+        if budgeted.included:
+            self._used_tokens += estimated
+            self._tier_usage[priority] = self._tier_usage.get(priority, 0) + estimated
+        else:
+            item_desc = f"{memory_type}:{getattr(item, 'id', 'unknown')}"
+            self._excluded.append(item_desc)
+
+        self._items.append(budgeted)
+        return budgeted
+
+    def apply_budget(
+        self,
+        memory_slice: MemorySlice,
+        type_priorities: Optional[Dict[str, PriorityTier]] = None,
+    ) -> Tuple[MemorySlice, BudgetReport]:
+        """
+        Apply budget constraints to a MemorySlice.
+
+        Args:
+            memory_slice: Raw retrieval results
+            type_priorities: Optional priority overrides per type
+
+        Returns:
+            Tuple of (budgeted MemorySlice, BudgetReport)
+        """
+        self.reset()
+
+        # Default priorities
+        priorities = type_priorities or {
+            "heuristic": PriorityTier.MUST_SEE,
+            "outcome": PriorityTier.SHOULD_SEE,
+            "domain_knowledge": PriorityTier.SHOULD_SEE,
+            "anti_pattern": PriorityTier.MUST_SEE,  # Important for avoiding mistakes
+            "preference": PriorityTier.MUST_SEE,  # User prefs are critical
+        }
+
+        # Process each type in priority order
+        included_heuristics = []
+        included_outcomes = []
+        included_knowledge = []
+        included_anti_patterns = []
+        included_preferences = []
+
+        # MUST_SEE first (preferences, anti-patterns, heuristics)
+        for pref in memory_slice.preferences[: self.config.max_preferences]:
+            budgeted = self.include(
+                pref, "preference", priorities.get("preference", PriorityTier.MUST_SEE)
+            )
+            if budgeted.included:
+                included_preferences.append(pref)
+
+        for ap in memory_slice.anti_patterns[: self.config.max_anti_patterns]:
+            budgeted = self.include(
+                ap,
+                "anti_pattern",
+                priorities.get("anti_pattern", PriorityTier.MUST_SEE),
+            )
+            if budgeted.included:
+                included_anti_patterns.append(ap)
+
+        for h in memory_slice.heuristics[: self.config.max_heuristics]:
+            budgeted = self.include(
+                h, "heuristic", priorities.get("heuristic", PriorityTier.MUST_SEE)
+            )
+            if budgeted.included:
+                included_heuristics.append(h)
+
+        # SHOULD_SEE next (outcomes, knowledge)
+        for o in memory_slice.outcomes[: self.config.max_outcomes]:
+            budgeted = self.include(
+                o, "outcome", priorities.get("outcome", PriorityTier.SHOULD_SEE)
+            )
+            if budgeted.included:
+                included_outcomes.append(o)
+
+        for k in memory_slice.domain_knowledge[: self.config.max_knowledge]:
+            budgeted = self.include(
+                k,
+                "domain_knowledge",
+                priorities.get("domain_knowledge", PriorityTier.SHOULD_SEE),
+            )
+            if budgeted.included:
+                included_knowledge.append(k)
+
+        # Build report
+        report = BudgetReport(
+            total_budget=self.config.max_tokens,
+            used_tokens=self._used_tokens,
+            remaining_tokens=self.remaining_tokens,
+            must_see_used=self._tier_usage.get(PriorityTier.MUST_SEE, 0),
+            must_see_budget=self.config.get_tier_budget(PriorityTier.MUST_SEE),
+            should_see_used=self._tier_usage.get(PriorityTier.SHOULD_SEE, 0),
+            should_see_budget=self.config.get_tier_budget(PriorityTier.SHOULD_SEE),
+            included_count=len([i for i in self._items if i.included]),
+            excluded_count=len([i for i in self._items if not i.included]),
+            truncated_count=len([i for i in self._items if i.truncated]),
+            summary_only_count=len([i for i in self._items if i.summary_only]),
+            budget_exceeded=self._used_tokens > self.config.max_tokens,
+            items_dropped=self._excluded,
+        )
+
+        # Build budgeted slice
+        budgeted_slice = MemorySlice(
+            heuristics=included_heuristics,
+            outcomes=included_outcomes,
+            preferences=included_preferences,
+            domain_knowledge=included_knowledge,
+            anti_patterns=included_anti_patterns,
+            query=memory_slice.query,
+            agent=memory_slice.agent,
+            retrieval_time_ms=memory_slice.retrieval_time_ms,
+        )
+
+        # Add budget metadata
+        budgeted_slice.metadata["budget_report"] = {
+            "total_budget": report.total_budget,
+            "used_tokens": report.used_tokens,
+            "utilization_pct": report.utilization_pct,
+            "items_dropped": len(report.items_dropped),
+        }
+
+        logger.info(
+            f"Budget applied: {report.used_tokens}/{report.total_budget} tokens "
+            f"({report.utilization_pct:.1f}%), "
+            f"{report.included_count} included, {report.excluded_count} excluded"
+        )
+
+        return budgeted_slice, report
+
+    def _default_classifier(self, item: Any, memory_type: str) -> PriorityTier:
+        """Default priority classification based on memory type and attributes."""
+        # Anti-patterns and preferences are always high priority
+        if memory_type in ("anti_pattern", "preference"):
+            return PriorityTier.MUST_SEE
+
+        # High-confidence heuristics are must-see
+        if memory_type == "heuristic":
+            if hasattr(item, "confidence") and item.confidence >= 0.8:
+                return PriorityTier.MUST_SEE
+            return PriorityTier.SHOULD_SEE
+
+        # Recent successful outcomes
+        if memory_type == "outcome":
+            if hasattr(item, "success") and item.success:
+                return PriorityTier.SHOULD_SEE
+            return PriorityTier.FETCH_ON_DEMAND
+
+        # Domain knowledge by confidence
+        if memory_type == "domain_knowledge":
+            if hasattr(item, "confidence") and item.confidence >= 0.7:
+                return PriorityTier.SHOULD_SEE
+            return PriorityTier.FETCH_ON_DEMAND
+
+        return PriorityTier.SHOULD_SEE
+
+    def get_fetch_on_demand_ids(self) -> List[str]:
+        """Get IDs of items marked for fetch-on-demand."""
+        return [
+            getattr(i.item, "id", None)
+            for i in self._items
+            if i.priority == PriorityTier.FETCH_ON_DEMAND and hasattr(i.item, "id")
+        ]
+
+
+class BudgetAwareRetrieval:
+    """
+    Wrapper that adds budget management to retrieval operations.
+
+    Usage:
+        budget_retrieval = BudgetAwareRetrieval(
+            retrieval_engine,
+            budget_config=BudgetConfig(max_tokens=4000)
+        )
+
+        result, report = budget_retrieval.retrieve_with_budget(
+            query="...",
+            agent="helena",
+            project_id="my-project"
+        )
+    """
+
+    def __init__(
+        self,
+        retrieval_engine: Any,  # RetrievalEngine
+        budget_config: Optional[BudgetConfig] = None,
+    ):
+        self.engine = retrieval_engine
+        self.budget = RetrievalBudget(config=budget_config)
+
+    def retrieve_with_budget(
+        self,
+        query: str,
+        agent: str,
+        project_id: str,
+        user_id: Optional[str] = None,
+        top_k: int = 10,
+        **kwargs,
+    ) -> Tuple[MemorySlice, BudgetReport]:
+        """
+        Retrieve memories with budget enforcement.
+
+        Returns:
+            Tuple of (budgeted MemorySlice, BudgetReport)
+        """
+        # Get raw results (request more than needed for budget filtering)
+        raw_slice = self.engine.retrieve(
+            query=query,
+            agent=agent,
+            project_id=project_id,
+            user_id=user_id,
+            top_k=top_k * 2,  # Get extra for filtering
+            **kwargs,
+        )
+
+        # Apply budget
+        return self.budget.apply_budget(raw_slice)