Upload aether/memory.py

aether/memory.py

@@ -0,0 +1,300 @@
+"""
+CoALA-inspired Memory Architecture for AETHER.
+Four modules: Working, Episodic, Semantic, Procedural.
+Plus Temporal Memory for long-horizon reasoning.
+"""
+
+import torch
+import torch.nn as nn
+from collections import deque
+from typing import Dict, List, Any, Optional
+import time
+import json
+
+
+class WorkingMemory:
+    """
+    Active scratchpad for current reasoning cycle.
+    Limited capacity, fast access.
+    """
+    def __init__(self, capacity: int = 16):
+        self.capacity = capacity
+        self.buffer: deque = deque(maxlen=capacity)
+        self.attention_weights = nn.Parameter(torch.ones(capacity))
+    
+    def store(self, item: Dict[str, Any]):
+        item["_timestamp"] = time.time()
+        self.buffer.append(item)
+    
+    def retrieve(self, query: str, top_k: int = 3) -> List[Dict]:
+        """Simple attention-based retrieval from working buffer."""
+        if not self.buffer:
+            return []
+        
+        # Compute relevance scores (simplified)
+        scores = []
+        for i, item in enumerate(self.buffer):
+            score = sum(1 for k in item if query.lower() in str(k).lower())
+            scores.append(score * torch.sigmoid(self.attention_weights[i]).item())
+        
+        # Get top-k indices
+        indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_k]
+        return [list(self.buffer)[i] for i in indices]
+    
+    def export(self) -> List[Dict]:
+        return list(self.buffer)
+    
+    def __len__(self):
+        return len(self.buffer)
+
+
+class EpisodicMemory:
+    """
+    Experience buffer storing past interactions.
+    Temporal structure for long-horizon reasoning.
+    """
+    def __init__(self, buffer_size: int = 1000):
+        self.buffer_size = buffer_size
+        self.buffer: deque = deque(maxlen=buffer_size)
+    
+    def store(self, episode: Dict[str, Any]):
+        episode["_timestamp"] = time.time()
+        self.buffer.append(episode)
+    
+    def retrieve_similar(self, query: str, top_k: int = 5) -> List[Dict]:
+        """Retrieve episodes similar to query."""
+        if not self.buffer:
+            return []
+        
+        # Simple keyword matching (replace with embedding-based in production)
+        scores = []
+        for item in self.buffer:
+            text = json.dumps(item)
+            score = sum(1 for word in query.lower().split() if word in text.lower())
+            scores.append(score)
+        
+        indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_k]
+        return [list(self.buffer)[i] for i in indices]
+    
+    def get_recent(self, n: int = 10) -> List[Dict]:
+        """Get n most recent episodes."""
+        return list(self.buffer)[-n:]
+    
+    def export(self) -> List[Dict]:
+        return list(self.buffer)
+    
+    def __len__(self):
+        return len(self.buffer)
+
+
+class SemanticMemory:
+    """
+    World knowledge - external and learned facts.
+    Backed by knowledge graph (see knowledge.py).
+    """
+    def __init__(self):
+        self.facts: Dict[str, Any] = {}
+    
+    def store_fact(self, key: str, value: Any, confidence: float = 1.0):
+        self.facts[key] = {"value": value, "confidence": confidence, "timestamp": time.time()}
+    
+    def retrieve(self, key: str) -> Optional[Dict]:
+        return self.facts.get(key)
+    
+    def query(self, query: str) -> List[Dict]:
+        """Simple prefix matching."""
+        return [v for k, v in self.facts.items() if query.lower() in k.lower()]
+    
+    def export(self) -> Dict:
+        return self.facts
+
+
+class ProceduralMemory:
+    """
+    Learned skills, tool definitions, and code implementations.
+    Inspired by Yunjue Agent's tool accumulation.
+    """
+    def __init__(self):
+        self.tools: Dict[str, Dict] = {}
+        self.tool_usage_stats: Dict[str, int] = {}
+    
+    def register_tool(self, name: str, code: str, description: str, 
+                     tags: List[str] = None):
+        self.tools[name] = {
+            "code": code,
+            "description": description,
+            "tags": tags or [],
+            "registered_at": time.time(),
+            "version": 1,
+        }
+        self.tool_usage_stats[name] = 0
+    
+    def get_tool(self, name: str) -> Optional[Dict]:
+        if name in self.tools:
+            self.tool_usage_stats[name] += 1
+            return self.tools[name]
+        return None
+    
+    def search_tools(self, query: str) -> List[Dict]:
+        """Search tools by description or tags."""
+        results = []
+        for name, tool in self.tools.items():
+            text = f"{name} {tool['description']} {' '.join(tool['tags'])}"
+            if query.lower() in text.lower():
+                results.append({"name": name, **tool})
+        return results
+    
+    def merge_tools(self, tool_cluster: List[str]) -> Optional[str]:
+        """
+        Merge functionally redundant tools (Yunjue-style tool absorption).
+        Returns name of merged tool or None.
+        """
+        if len(tool_cluster) < 2:
+            return None
+        
+        # Simple merge: keep highest usage tool as canonical
+        canonical = max(tool_cluster, key=lambda t: self.tool_usage_stats.get(t, 0))
+        
+        # Merge descriptions
+        merged_desc = " | ".join(
+            self.tools[t]["description"] for t in tool_cluster if t in self.tools
+        )
+        self.tools[canonical]["description"] = merged_desc
+        self.tools[canonical]["version"] += 1
+        
+        # Remove redundant tools
+        for t in tool_cluster:
+            if t != canonical and t in self.tools:
+                del self.tools[t]
+        
+        return canonical
+    
+    def export(self) -> Dict:
+        return {
+            "tools": self.tools,
+            "usage_stats": self.tool_usage_stats,
+        }
+
+
+class CoALAMemory:
+    """
+    Unified memory system following CoALA cognitive architecture.
+    Combines Working, Episodic, Semantic, and Procedural memory.
+    """
+    def __init__(self, capacity: int = 16):
+        self.working = WorkingMemory(capacity=capacity)
+        self.episodic = EpisodicMemory(buffer_size=1000)
+        self.semantic = SemanticMemory()
+        self.procedural = ProceduralMemory()
+    
+    def store(self, item: Dict[str, Any], memory_type: str = "working"):
+        if memory_type == "working":
+            self.working.store(item)
+        elif memory_type == "episodic":
+            self.episodic.store(item)
+        elif memory_type == "semantic":
+            for k, v in item.items():
+                self.semantic.store_fact(k, v)
+        elif memory_type == "procedural":
+            if "name" in item and "code" in item:
+                self.procedural.register_tool(
+                    item["name"], item["code"], 
+                    item.get("description", ""),
+                    item.get("tags", [])
+                )
+    
+    def retrieve(self, query: str, memory_type: str = "all", top_k: int = 5) -> List[Dict]:
+        if memory_type == "all":
+            results = []
+            results.extend(self.working.retrieve(query, top_k=top_k//2))
+            results.extend(self.episodic.retrieve_similar(query, top_k=top_k))
+            results.extend(self.semantic.query(query)[:top_k])
+            return results[:top_k]
+        elif memory_type == "working":
+            return self.working.retrieve(query, top_k)
+        elif memory_type == "episodic":
+            return self.episodic.retrieve_similar(query, top_k)
+        elif memory_type == "semantic":
+            return self.semantic.query(query)[:top_k]
+        elif memory_type == "procedural":
+            return self.procedural.search_tools(query)
+        return []
+    
+    @property
+    def buffer(self):
+        """Alias for working memory buffer."""
+        return self.working.buffer
+    
+    def export(self) -> Dict[str, Any]:
+        return {
+            "working": self.working.export(),
+            "episodic": self.episodic.export(),
+            "semantic": self.semantic.export(),
+            "procedural": self.procedural.export(),
+        }
+
+
+class TemporalMemory(nn.Module):
+    """
+    Time-sensitive memory with learned temporal attention.
+    Enables long-horizon reasoning and contextual adaptation.
+    Uses a simple LSTM-like gating mechanism.
+    """
+    def __init__(self, buffer_size: int = 1000, hidden_dim: int = 64):
+        super().__init__()
+        self.buffer_size = buffer_size
+        self.hidden_dim = hidden_dim
+        self.buffer: deque = deque(maxlen=buffer_size)
+        
+        # Temporal attention network
+        self.temporal_gate = nn.Sequential(
+            nn.Linear(2, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 1),
+            nn.Sigmoid(),
+        )
+    
+    def store(self, event: Dict[str, Any]):
+        event["_timestamp"] = time.time()
+        self.buffer.append(event)
+    
+    def retrieve_context(self, current_time: Optional[float] = None,
+                         lookback_window: float = 3600.0) -> List[Dict]:
+        """
+        Retrieve events within lookback window, weighted by recency.
+        """
+        current_time = current_time or time.time()
+        relevant = []
+        
+        for event in self.buffer:
+            age = current_time - event.get("_timestamp", current_time)
+            if age <= lookback_window:
+                # Temporal relevance score: exponential decay
+                recency_score = torch.exp(torch.tensor(-age / lookback_window))
+                relevant.append({
+                    **event,
+                    "recency_score": recency_score.item(),
+                    "age_seconds": age,
+                })
+        
+        # Sort by recency score
+        relevant.sort(key=lambda x: x["recency_score"], reverse=True)
+        return relevant
+    
+    def retrieve_with_attention(self, query_embedding: torch.Tensor,
+                                   top_k: int = 10) -> List[Dict]:
+        """
+        Attention-based retrieval combining temporal and semantic relevance.
+        (Placeholder - would use actual embeddings in full implementation)
+        """
+        return self.retrieve_context()[:top_k]
+    
+    def export(self) -> List[Dict]:
+        return list(self.buffer)
+    
+    @property
+    def buffer_contents(self):
+        return list(self.buffer)
+    
+    def __len__(self):
+        return len(self.buffer)