""" Attention-Weighted Context Window Pruning Module Implements hierarchical attention scoring for context window optimization, based on predictive coding (Rao & Ballard, 1999), episodic buffer management (Baddeley, 2000), and KV cache compression (Xiong et al., 2023). Usage: from context_window_attention_model import ContextItem, AttentionScorer scorer = AttentionScorer() scorer.add_item(ContextItem( content="Current task directive", item_type="system_prompt", age_in_cycles=0, reference_count=3, task_relevance=1.0 )) scored_items = scorer.score_all() compressed = scorer.compress(scored_items, threshold=0.25) """ import math from dataclasses import dataclass, field from enum import Enum from typing import Optional class ContextTier(Enum): """Hierarchical context tiers based on cognitive role.""" PFC = "pfc" HIPPOCAMPUS = "hippocampus" SENSORY = "sensory" class ItemType(Enum): """Context item types with base attention weights.""" SYSTEM_PROMPT = "system_prompt" CURRENT_TASK = "current_task" STEERING_DIRECTIVE = "steering_directive" CHAT_HISTORY = "chat_history" RELATIONSHIP_STATE = "relationship_state" ACTIVE_MEMORY = "active_memory" MEDIA_TRANSCRIPTION = "media_transcription" TOOL_OUTPUT = "tool_output" STALE_METADATA = "stale_metadata" @property def base_weight(self) -> float: weights = { ItemType.SYSTEM_PROMPT: 1.0, ItemType.CURRENT_TASK: 1.0, ItemType.STEERING_DIRECTIVE: 1.0, ItemType.CHAT_HISTORY: 0.7, ItemType.RELATIONSHIP_STATE: 0.6, ItemType.ACTIVE_MEMORY: 0.5, ItemType.MEDIA_TRANSCRIPTION: 0.3, ItemType.TOOL_OUTPUT: 0.3, ItemType.STALE_METADATA: 0.2, } return weights[self] @property def tier(self) -> ContextTier: tier_map = { ItemType.SYSTEM_PROMPT: ContextTier.PFC, ItemType.CURRENT_TASK: ContextTier.PFC, ItemType.STEERING_DIRECTIVE: ContextTier.PFC, ItemType.CHAT_HISTORY: ContextTier.HIPPOCAMPUS, ItemType.RELATIONSHIP_STATE: ContextTier.HIPPOCAMPUS, ItemType.ACTIVE_MEMORY: ContextTier.HIPPOCAMPUS, ItemType.MEDIA_TRANSCRIPTION: ContextTier.SENSORY, ItemType.TOOL_OUTPUT: ContextTier.SENSORY, ItemType.STALE_METADATA: ContextTier.SENSORY, } return tier_map[self] @dataclass class ContextItem: """A single context item with attention metadata.""" content: str item_type: ItemType age_in_cycles: float = 0.0 reference_count: int = 0 task_relevance: float = 0.5 is_active: bool = True metadata: dict = field(default_factory=dict) def __post_init__(self): self.reference_count = max(0, self.reference_count) self.task_relevance = max(0.0, min(1.0, self.task_relevance)) self.age_in_cycles = max(0.0, self.age_in_cycles) @dataclass class ScoredItem: """A context item with its computed attention score.""" item: ContextItem attention_score: float tier: ContextTier is_compressed: bool = False @property def token_estimate(self) -> int: return max(1, len(self.item.content) // 4) class AttentionScorer: """ Computes attention scores for context items using the formula: attention_score(item) = base_weight * recency_factor * reference_count * task_relevance Where: base_weight: Item type's inherent importance (0.2-1.0) recency_factor: exp(-lambda * age_in_cycles), lambda=0.3 reference_count: min(1 + 0.1 * refs, 1.5) -- capped at +0.5 bonus task_relevance: 0.25 (unrelated) to 1.0 (directly related) """ def __init__(self, recency_lambda: float = 0.3, ref_bonus: float = 0.1, ref_cap: float = 0.5): self.recency_lambda = recency_lambda self.ref_bonus = ref_bonus self.ref_cap = ref_cap self.items: list[ContextItem] = [] def add_item(self, item: ContextItem) -> None: self.items.append(item) def add_items(self, items: list[ContextItem]) -> None: self.items.extend(items) def compute_recency_factor(self, age_in_cycles: float) -> float: return math.exp(-self.recency_lambda * age_in_cycles) def compute_reference_bonus(self, reference_count: int) -> float: return min(self.ref_bonus * reference_count, self.ref_cap) def score_item(self, item: ContextItem) -> float: base = item.item_type.base_weight recency = self.compute_recency_factor(item.age_in_cycles) ref_bonus = self.compute_reference_bonus(item.reference_count) relevance = item.task_relevance score = base * recency * (1.0 + ref_bonus) * relevance return round(score, 4) def score_all(self) -> list[ScoredItem]: scored = [] for item in self.items: score = self.score_item(item) scored_item = ScoredItem( item=item, attention_score=score, tier=item.item_type.tier, ) scored.append(scored_item) tier_priority = { ContextTier.PFC: 0, ContextTier.HIPPOCAMPUS: 1, ContextTier.SENSORY: 2, } scored.sort(key=lambda x: (-x.attention_score, tier_priority[x.tier])) return scored def compress(self, scored_items: list[ScoredItem], threshold: float = 0.25, max_tokens: int = 10000) -> list[ScoredItem]: compressed = [] total_tokens = 0 for item in scored_items: if item.is_compressed: continue if item.attention_score < threshold: if item.item.age_in_cycles > 10 and item.item.reference_count == 0: item.is_compressed = True continue item_tokens = item.token_estimate if total_tokens + item_tokens > max_tokens: item.item.content = item.item.content[:len(item.item.content) // 2] + ".." item.is_compressed = True continue compressed.append(item) total_tokens += item_tokens return compressed def get_summary(self, scored_items: list[ScoredItem]) -> dict: total_tokens = sum(item.token_estimate for item in scored_items) tier_counts = {} for item in scored_items: tier = item.tier.value tier_counts[tier] = tier_counts.get(tier, 0) + 1 return { "total_items": len(scored_items), "total_tokens": total_tokens, "tier_distribution": tier_counts, "avg_score": sum(i.attention_score for i in scored_items) / len(scored_items) if scored_items else 0, "compressed_count": sum(1 for i in scored_items if i.is_compressed), } def build_default_context() -> list[ContextItem]: items = [ ContextItem( content="You are replying to the authenticated Telegram admin in a private DM.", item_type=ItemType.SYSTEM_PROMPT, age_in_cycles=0, reference_count=5, task_relevance=1.0, ), ContextItem( content="Current task: implement attention-weighted pruning for context window.", item_type=ItemType.CURRENT_TASK, age_in_cycles=0, reference_count=3, task_relevance=1.0, ), ContextItem( content="Telegram live context from @robitman: tell me concretely with research paper references.", item_type=ItemType.STEERING_DIRECTIVE, age_in_cycles=1, reference_count=2, task_relevance=1.0, ), ContextItem( content="11:35 PM @robitman/chat: text='tell me concretely with research paper references exactly how you would like me to update this'", item_type=ItemType.CHAT_HISTORY, age_in_cycles=0, reference_count=1, task_relevance=0.8, ), ContextItem( content="Relationship: @robitman --affirmed--> @omnius_agent_bot confidence=0.52 weight=1.00", item_type=ItemType.RELATIONSHIP_STATE, age_in_cycles=2, reference_count=1, task_relevance=0.6, ), ContextItem( content="Active memory: @robitman shared media: voice, audio/ogg, 10s, 43595 bytes", item_type=ItemType.ACTIVE_MEMORY, age_in_cycles=3, reference_count=1, task_relevance=0.5, ), ContextItem( content="[voice message transcribed: 'So you're suggesting omitting large swaths of elements of the context window that are just simply not relevant to the current context at hand.']", item_type=ItemType.MEDIA_TRANSCRIPTION, age_in_cycles=3, reference_count=1, task_relevance=0.4, ), ContextItem( content="[voice message transcribed: 'How should I restructure your context window taking into account attention mechanisms?']", item_type=ItemType.MEDIA_TRANSCRIPTION, age_in_cycles=4, reference_count=0, task_relevance=0.3, ), ContextItem( content="[voice message transcribed: 'How do your procedures appear to you currently and what should we implement on Omnius, the coding agent, and play here to help prevent these failures in the future and guarantee better critiques of yourself and your actions?']", item_type=ItemType.MEDIA_TRANSCRIPTION, age_in_cycles=5, reference_count=0, task_relevance=0.3, ), ContextItem( content="[voice message transcribed: 'In the future, when we run into issues where you take an action and the actions resulted in failures, how do we account for these failures in a way where you check your work before deeming success or check your work before deeming failure when you may have succeeded leading to duplicates?']", item_type=ItemType.MEDIA_TRANSCRIPTION, age_in_cycles=6, reference_count=0, task_relevance=0.2, ), ContextItem( content="[voice message transcribed: 'You created like four duplicates.']", item_type=ItemType.MEDIA_TRANSCRIPTION, age_in_cycles=7, reference_count=0, task_relevance=0.2, ), ContextItem( content="[voice message transcribed: 'Delete all of the duplock kits.']", item_type=ItemType.MEDIA_TRANSCRIPTION, age_in_cycles=8, reference_count=0, task_relevance=0.2, ), ContextItem( content="[MID_TASK_STEERING_INTAKE v2] Source: injected user message during an active run.", item_type=ItemType.TOOL_OUTPUT, age_in_cycles=5, reference_count=0, task_relevance=0.3, ), ContextItem( content="[SYSTEM] You have 3 failed approaches this session. Consider using memory_write to save these failure patterns.", item_type=ItemType.TOOL_OUTPUT, age_in_cycles=6, reference_count=0, task_relevance=0.3, ), ContextItem( content="[PROGRESS GATE - evidence gathered, no files changed] Successful discovery calls: 3.", item_type=ItemType.TOOL_OUTPUT, age_in_cycles=7, reference_count=0, task_relevance=0.2, ), ContextItem( content="[REG-61 directive active] A REG-61 FIRST-EDIT NUDGE was issued earlier and has not yet been satisfied.", item_type=ItemType.TOOL_OUTPUT, age_in_cycles=8, reference_count=0, task_relevance=0.2, ), ContextItem( content="[STOP - RETRY LOOP DETECTED] You are re-issuing the SAME failing tool call(s) without changing anything.", item_type=ItemType.TOOL_OUTPUT, age_in_cycles=9, reference_count=0, task_relevance=0.2, ), ContextItem( content="[world-state turn=8] GOAL: You are replying to the authenticated Telegram admin in a private DM.", item_type=ItemType.TOOL_OUTPUT, age_in_cycles=10, reference_count=0, task_relevance=0.15, ), ContextItem( content="[RECENT UNRESOLVED FAILURES] file_write:content=# Context Window Optimization Spec attempts=2", item_type=ItemType.TOOL_OUTPUT, age_in_cycles=10, reference_count=0, task_relevance=0.15, ), ContextItem( content="[SHELL FAILURE PIVOT - raw output repeated] Recent failed shell calls: 2.", item_type=ItemType.TOOL_OUTPUT, age_in_cycles=10, reference_count=0, task_relevance=0.15, ), ContextItem( content="[TRIED: file_read, file_write, list_directory, shell] No creative edits yet this run.", item_type=ItemType.STALE_METADATA, age_in_cycles=12, reference_count=0, task_relevance=0.1, ), ContextItem( content="[TRIED: find . -name '*.py' -path '*/bridge*' -o -name '*.py' -path '*/context*']", item_type=ItemType.STALE_METADATA, age_in_cycles=12, reference_count=0, task_relevance=0.1, ), ] return items def demonstrate_optimization() -> dict: scorer = AttentionScorer() items = build_default_context() scorer.add_items(items) scored = scorer.score_all() print("=" * 60) print("ATTENTION-WEIGHTED CONTEXT WINDOW OPTIMIZATION") print("=" * 60) print() print(f"{'Item':<50} {'Score':>6} {'Tier':<12} {'Tokens':>6}") print("-" * 75) for item in scored: content_preview = item.item.content[:48].replace('\n', ' ') print(f"{content_preview:<50} {item.attention_score:>6.4f} {item.tier.value:<12} {item.token_estimate:>6}") print() print(f"Total items: {len(scored)}") total_tokens = sum(i.token_estimate for i in scored) print(f"Total tokens: {total_tokens}") print(f"Avg score: {sum(i.attention_score for i in scored) / len(scored):.4f}") compressed = scorer.compress(scored, threshold=0.25, max_tokens=10000) compressed_tokens = sum(i.token_estimate for i in compressed) reduction = (1 - compressed_tokens / total_tokens) * 100 if total_tokens > 0 else 0 print() print(f"After compression (threshold=0.25):") print(f" Items retained: {len(compressed)}/{len(scored)}") print(f" Tokens: {compressed_tokens}/{total_tokens}") print(f" Reduction: {reduction:.1f}%") return { "total_items": len(scored), "total_tokens": total_tokens, "compressed_items": len(compressed), "compressed_tokens": compressed_tokens, "reduction_pct": round(reduction, 1), "avg_score": round(sum(i.attention_score for i in scored) / len(scored), 4), } if __name__ == "__main__": result = demonstrate_optimization() print() print(f"Result: {result}")