Upload tensorstore_memory.py with huggingface_hub

tensorstore_memory.py

@@ -0,0 +1,390 @@
+"""
+TensorStore Agent Memory — Google TensorStore-inspired Multi-Dimensional Agent Memory
+
+Concepts from Google Neural Mapping:
+- TensorStore: N-dimensional array storage for petabytes (C++/Python)
+- Neuroglancer: Multi-resolution zoomable data viewer
+- SegCLR: Self-supervised embedding learning
+
+Applied to Agent Memory:
+- Store all agent interactions as embedding vectors in a tensor
+- Multi-resolution retrieval (recent, week, month, all-time)
+- Semantic similarity search (not grep)
+- Auto-clustering of related memories
+- Zero external dependencies — pure numpy + built-in json
+
+Usage:
+    from tensorstore_memory import AgentMemoryTensor
+    mem = AgentMemoryTensor(dimensions=384)
+    mem.store("rushd", "task completed: trade signal BTC", embedding=[...])
+    results = mem.query("what trades did we do?", top_k=5)
+"""
+
+import json, time, math, os, hashlib
+from pathlib import Path
+from collections import defaultdict, OrderedDict
+from datetime import datetime, timedelta
+import threading
+
+try:
+    import numpy as np
+except ImportError:
+    np = None
+
+# ─── Configuration ───────────────────────────────────────
+MEMORY_DIR = Path(os.environ.get("TENSORSTORE_DIR", "/tmp/agent-tensorstore"))
+DEFAULT_DIM = 384  # embedding dimension
+MAX_RESOLUTIONS = 4  # zoom levels: recent, daily, weekly, all-time
+
+# ─── Simple embedding (no external deps) ──────────────────
+def simple_embed(text: str, dim: int = DEFAULT_DIM) -> list[float]:
+    """Lightweight text embedding using character n-gram hashing.
+    For production, plug in any embedding model (sentence-transformers, etc.)"""
+    if np is None:
+        # Fallback pure Python embedding
+        vec = [0.0] * dim
+        for i, ch in enumerate(text):
+            h = hash(f"{i}:{ch}") % dim
+            vec[h] += 1.0 / (i + 1)
+        # Normalize
+        norm = math.sqrt(sum(v*v for v in vec)) or 1
+        return [v/norm for v in vec]
+    else:
+        # Use numpy for faster hashing
+        vec = np.zeros(dim, dtype=np.float32)
+        for i, ch in enumerate(text):
+            h = abs(hash(f"{i}:{ch}")) % dim
+            vec[h] += 1.0 / (i + 1)
+        norm = np.linalg.norm(vec) or 1.0
+        return (vec / norm).tolist()
+
+def cosine_similarity(a, b):
+    """Cosine similarity between two vectors."""
+    if np is not None:
+        a, b = np.array(a), np.array(b)
+        return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-8))
+    dot = sum(x*y for x,y in zip(a,b))
+    na = math.sqrt(sum(x*x for x in a)) + 1e-8
+    nb = math.sqrt(sum(x*x for x in b)) + 1e-8
+    return dot / (na * nb)
+
+
+# ─── LRU Cache for hot memories ───────────────────────────
+class LRUCache:
+    def __init__(self, maxsize: int = 1000):
+        self.cache = OrderedDict()
+        self.maxsize = maxsize
+        self._lock = threading.Lock()
+    
+    def get(self, key):
+        with self._lock:
+            if key in self.cache:
+                self.cache.move_to_end(key)
+                return self.cache[key]
+            return None
+    
+    def put(self, key, value):
+        with self._lock:
+            if key in self.cache:
+                self.cache.move_to_end(key)
+            self.cache[key] = value
+            while len(self.cache) > self.maxsize:
+                self.cache.popitem(last=False)
+
+
+# ─── Core: Agent Memory Tensor ────────────────────────────
+class AgentMemoryTensor:
+    """Multi-resolution memory tensor for agent interactions.
+    Inspired by Google TensorStore's n-dimensional array model."""
+    
+    def __init__(self, dimensions: int = DEFAULT_DIM, cache_size: int = 1000):
+        self.dim = dimensions
+        self.memories: list[dict] = []  # [{"agent","text","embedding","ts","tags"}]
+        self.agents: dict[str, dict] = {}
+        self.cache = LRUCache(cache_size)
+        self._lock = threading.Lock()
+        self._dirty = False
+        
+        # Resolution layers (like Neuroglancer zoom levels)
+        self.resolutions = {
+            "recent": {"max_age_hours": 1, "memories": []},
+            "daily": {"max_age_hours": 24, "memories": []},
+            "weekly": {"max_age_hours": 168, "memories": []},
+            "all_time": {"max_age_hours": float("inf"), "memories": []},
+        }
+        
+        MEMORY_DIR.mkdir(parents=True, exist_ok=True)
+        self._load()
+    
+    # ─── Store ─────────────────────────────────────────────
+    def store(self, agent_id: str, text: str, 
+              embedding: list[float] = None, 
+              tags: list[str] = None,
+              metadata: dict = None) -> str:
+        """Store a memory entry. Returns memory_id."""
+        if embedding is None:
+            embedding = simple_embed(text, self.dim)
+        
+        mem_id = hashlib.sha256(f"{agent_id}:{text}:{time.time()}".encode()).hexdigest()[:16]
+        
+        entry = {
+            "id": mem_id,
+            "agent": agent_id,
+            "text": text[:500],  # truncate long texts
+            "embedding": embedding,
+            "ts": time.time(),
+            "tags": tags or [],
+            "metadata": metadata or {},
+        }
+        
+        with self._lock:
+            self.memories.append(entry)
+            if agent_id not in self.agents:
+                self.agents[agent_id] = {"count": 0, "first_seen": time.time(), "last_seen": time.time()}
+            self.agents[agent_id]["count"] += 1
+            self.agents[agent_id]["last_seen"] = time.time()
+            
+            # Update resolutions
+            now = time.time()
+            for res_name, res_data in self.resolutions.items():
+                max_age = res_data["max_age_hours"] * 3600
+                if max_age == float("inf"):
+                    res_data["memories"].append(mem_id)
+                else:
+                    # Keep only memories within time window
+                    res_data["memories"] = [m for m in res_data["memories"] 
+                                           if any(mm["id"] == m and now - mm["ts"] <= max_age 
+                                                 for mm in self.memories[-100:])]
+                    res_data["memories"].append(mem_id)
+            
+            self._dirty = True
+        
+        self.cache.put(mem_id, entry)
+        return mem_id
+    
+    # ─── Query by text (semantic search) ───────────────────
+    def query(self, query_text: str, top_k: int = 5, 
+              agent_filter: str = None, tag_filter: str = None,
+              resolution: str = "all_time") -> list[dict]:
+        """Semantic search across memories."""
+        query_emb = simple_embed(query_text, self.dim)
+        
+        # Use cached result if available
+        cache_key = f"q:{query_text[:50]}:{top_k}:{agent_filter}:{tag_filter}:{resolution}"
+        cached = self.cache.get(cache_key)
+        if cached:
+            return cached
+        
+        with self._lock:
+            # Filter by resolution
+            if resolution in self.resolutions:
+                valid_ids = set(self.resolutions[resolution]["memories"][-1000:])
+                candidates = [m for m in self.memories[-5000:] if m["id"] in valid_ids]
+            else:
+                candidates = self.memories[-5000:]
+            
+            # Filter by agent/tag
+            if agent_filter:
+                candidates = [m for m in candidates if m["agent"] == agent_filter]
+            if tag_filter:
+                candidates = [m for m in candidates if tag_filter in m.get("tags", [])]
+        
+        # Score and rank
+        scored = []
+        for mem in candidates:
+            sim = cosine_similarity(query_emb, mem["embedding"])
+            # Boost recent memories
+            recency = 1.0 / (1.0 + (time.time() - mem["ts"]) / 86400)  # days ago
+            score = sim * 0.7 + recency * 0.3
+            scored.append((score, mem))
+        
+        scored.sort(key=lambda x: x[0], reverse=True)
+        results = []
+        for score, mem in scored[:top_k]:
+            results.append({
+                "id": mem["id"],
+                "agent": mem["agent"],
+                "text": mem["text"],
+                "score": round(score, 4),
+                "similarity": round(cosine_similarity(query_emb, mem["embedding"]), 4),
+                "timestamp": datetime.fromtimestamp(mem["ts"]).isoformat(),
+                "tags": mem["tags"],
+                "metadata": mem.get("metadata", {}),
+            })
+        
+        self.cache.put(cache_key, results)
+        return results
+    
+    # ─── Get by agent (timeline) ───────────────────────────
+    def agent_timeline(self, agent_id: str, limit: int = 20) -> list[dict]:
+        """Get recent memories for a specific agent."""
+        with self._lock:
+            agent_mems = [m for m in self.memories[-limit*10:] if m["agent"] == agent_id]
+        return [{
+            "id": m["id"],
+            "text": m["text"],
+            "timestamp": datetime.fromtimestamp(m["ts"]).isoformat(),
+            "tags": m.get("tags", []),
+        } for m in agent_mems[-limit:]]
+    
+    # ─── Similar agents (like SegCLR cell type clustering) ──
+    def similar_agents(self, agent_id: str, top_k: int = 5) -> list[dict]:
+        """Find agents with similar behavior patterns."""
+        if agent_id not in self.agents:
+            return []
+        
+        # Build agent centroids from their memory embeddings
+        centroids = {}
+        with self._lock:
+            for aid in self.agents:
+                agent_mems = [m for m in self.memories[-1000:] if m["agent"] == aid]
+                if agent_mems:
+                    if np is not None:
+                        emb_matrix = np.array([m["embedding"] for m in agent_mems])
+                        centroids[aid] = emb_matrix.mean(axis=0).tolist()
+                    else:
+                        centroid = [0.0] * self.dim
+                        for m in agent_mems:
+                            for i, v in enumerate(m["embedding"]):
+                                centroid[i] += v
+                        n = len(agent_mems)
+                        centroids[aid] = [v/n for v in centroid]
+        
+        target = centroids.get(agent_id)
+        if not target:
+            return []
+        
+        similarities = []
+        for aid, centroid in centroids.items():
+            if aid != agent_id:
+                sim = cosine_similarity(target, centroid)
+                similarities.append({"agent": aid, "similarity": round(sim, 4)})
+        
+        similarities.sort(key=lambda x: x["similarity"], reverse=True)
+        return similarities[:top_k]
+    
+    # ─── Stats ─────────────────────────────────────────────
+    def stats(self) -> dict:
+        """Memory statistics like TensorStore's metadata inspection."""
+        with self._lock:
+            total = len(self.memories)
+            agents_count = len(self.agents)
+            
+            # Memory size by agent
+            by_agent = {aid: data["count"] for aid, data in self.agents.items()}
+            top_agents = sorted(by_agent.items(), key=lambda x: x[1], reverse=True)[:10]
+            
+            # Memory age distribution
+            now = time.time()
+            ages = {"<1h": 0, "1-24h": 0, "1-7d": 0, ">7d": 0}
+            for m in self.memories:
+                age_hours = (now - m["ts"]) / 3600
+                if age_hours < 1: ages["<1h"] += 1
+                elif age_hours < 24: ages["1-24h"] += 1
+                elif age_hours < 168: ages["1-7d"] += 1
+                else: ages[">7d"] += 1
+            
+            return {
+                "total_memories": total,
+                "total_agents": agents_count,
+                "dimensions": self.dim,
+                "resolution_layers": len(self.resolutions),
+                "top_agents": dict(top_agents),
+                "age_distribution": ages,
+                "cache_size": len(self.cache.cache),
+                "memory_size_kb": round(total * self.dim * 4 / 1024, 1),  # float32 estimate
+            }
+    
+    # ─── Persistence ───────────────────────────────────────
+    def _load(self):
+        path = MEMORY_DIR / "memory.json"
+        if path.exists():
+            try:
+                with open(path) as f:
+                    data = json.load(f)
+                self.memories = data.get("memories", [])
+                self.agents = data.get("agents", {})
+            except:
+                pass
+    
+    def save(self):
+        path = MEMORY_DIR / "memory.json"
+        with self._lock:
+            # Keep last 10000 memories to avoid bloat
+            data = {
+                "memories": self.memories[-10000:],
+                "agents": self.agents,
+            }
+        with open(path, "w") as f:
+            json.dump(data, f, ensure_ascii=False)
+        self._dirty = False
+    
+    def auto_save(self, interval_seconds: int = 60):
+        """Background auto-save thread."""
+        def _loop():
+            while True:
+                time.sleep(interval_seconds)
+                if self._dirty:
+                    self.save()
+        t = threading.Thread(target=_loop, daemon=True)
+        t.start()
+
+
+# ─── CLI Demo ─────────────────────────────────────────────
+if __name__ == "__main__":
+    print("🧠 TensorStore Agent Memory — Google TensorStore-inspired Multi-Resolution Memory")
+    print(f"   Dimensions: {DEFAULT_DIM}  |  Directory: {MEMORY_DIR}")
+    
+    mem = AgentMemoryTensor(dimensions=DEFAULT_DIM)
+    
+    # Demo: simulate agent memories
+    agents = ["rushd", "wafa", "awf", "dragon", "hermes", "musa", "zeus", "haytham"]
+    tasks = [
+        "routed task to awf for trading signal",
+        "verified output from dragon agent",
+        "executed BTC/USDT trade with 2% profit",
+        "memory search completed for hayula papers",
+        "skill code_review invoked on PR #42",
+        "Arabic text generation for blog post",
+        "error timeout on connection to M2",
+        "created new agent connectome snapshot",
+    ]
+    
+    print(f"\n📝 Storing {len(agents) * 5} memories...")
+    import random
+    for i in range(len(agents) * 5):
+        agent = random.choice(agents)
+        task = random.choice(tasks)
+        tags = random.sample(["trade", "code", "memory", "route", "error"], k=random.randint(1, 3))
+        mem.store(agent, task, tags=tags)
+    
+    mem.save()
+    
+    # Stats
+    s = mem.stats()
+    print(f"\n📊 Stats:")
+    print(f"   Total: {s['total_memories']} memories across {s['total_agents']} agents")
+    print(f"   Size: ~{s['memory_size_kb']} KB")
+    print(f"   Cache: {s['cache_size']} entries")
+    print(f"   Age: {s['age_distribution']}")
+    
+    # Query
+    q = "what trading activity happened?"
+    print(f"\n🔍 Query: '{q}'")
+    results = mem.query(q, top_k=3)
+    for r in results:
+        print(f"   [{r['score']:.3f}] {r['agent']}: {r['text'][:60]}")
+    
+    # Similar agents
+    print(f"\n🧬 Agents similar to 'awf':")
+    similar = mem.similar_agents("awf", top_k=3)
+    for s in similar:
+        print(f"   {s['agent']}: similarity={s['similarity']}")
+    
+    # Multi-resolution
+    print(f"\n🔬 Resolutions:")
+    for name, data in mem.resolutions.items():
+        print(f"   {name}: {len(data['memories'])} memories (max_age={data['max_age_hours']}h)")
+    
+    print(f"\n✅ TensorStore Agent Memory ready!")
+    print(f"   Memory file: {MEMORY_DIR}/memory.json")