tensorstore_memory.py

"""
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")