CHKIM79/scalable-ai-agent-system

1"""2Multi-layered Memory Management System3Implements short-term, long-term, working, and episodic memory4"""5import asyncio6import json7import sqlite38import time9from typing import Dict, List, Any, Optional, Tuple10from dataclasses import dataclass, asdict11from datetime import datetime, timedelta12from enum import Enum13import numpy as np14import hashlib15# import chromadb16# from chromadb.config import Settings17 18 19class MemoryType(Enum):20    SHORT_TERM = "short_term"21    LONG_TERM = "long_term"22    WORKING = "working"23    EPISODIC = "episodic"24 25 26@dataclass27class MemoryEntry:28    id: str29    content: str30    memory_type: MemoryType31    timestamp: datetime32    importance: float = 0.533    access_count: int = 034    last_accessed: datetime = None35    metadata: Dict[str, Any] = None36    embedding: Optional[List[float]] = None37 38 39class MemoryManager:40    """41    Comprehensive memory management system with multiple memory types42    """43    44    def __init__(self, db_path: str = "memory.db", vector_db_path: str = "./vector_db"):45        self.db_path = db_path46        self.vector_db_path = vector_db_path47        48        # Simple hash-based embedding for testing49        self.embedding_model = None50        51        # Simplified vector storage52        self.vector_storage = {}53        54        # Memory stores55        self.short_term_memory: Dict[str, MemoryEntry] = {}56        self.working_memory: Dict[str, Any] = {}57        58        # TTL for short-term memory (in seconds)59        self.short_term_ttl = 3600  # 1 hour60        61        # Vector database for long-term memory62        self.vector_client = None63        self.vector_collection = None64        65        # SQLite for structured memory storage66        self.db_connection = None67        68    async def initialize(self):69        """Initialize all memory components"""70        # Initialize SQLite database71        self.db_connection = sqlite3.connect(self.db_path, check_same_thread=False)72        await self._create_tables()73        74        # Initialize simplified vector storage75        self.vector_storage = {}76        self.vector_collection = "simplified_storage"77        78        # Start cleanup task for short-term memory79        asyncio.create_task(self._cleanup_short_term_memory())80        81    async def _create_tables(self):82        """Create necessary database tables"""83        cursor = self.db_connection.cursor()84        85        # Episodic memory table86        cursor.execute('''87            CREATE TABLE IF NOT EXISTS episodic_memory (88                id TEXT PRIMARY KEY,89                content TEXT NOT NULL,90                timestamp DATETIME NOT NULL,91                importance REAL DEFAULT 0.5,92                access_count INTEGER DEFAULT 0,93                last_accessed DATETIME,94                metadata TEXT,95                embedding BLOB96            )97        ''')98        99        # Interaction history table100        cursor.execute('''101            CREATE TABLE IF NOT EXISTS interactions (102                id TEXT PRIMARY KEY,103                interaction_type TEXT NOT NULL,104                content TEXT NOT NULL,105                timestamp DATETIME NOT NULL,106                metadata TEXT107            )108        ''')109        110        self.db_connection.commit()111    112    async def store_memory(self, content: str, memory_type: str = "general", metadata: Dict[str, Any] = None):113        """Store a memory entry (unified API)"""114        return await self.store_interaction(memory_type, content, metadata)115    116    async def store_interaction(self, interaction_type: str, content: str, metadata: Dict[str, Any] = None):117        """Store an interaction in memory"""118        entry_id = f"{interaction_type}_{int(time.time() * 1000)}"119        120        # Store in short-term memory121        memory_entry = MemoryEntry(122            id=entry_id,123            content=content,124            memory_type=MemoryType.SHORT_TERM,125            timestamp=datetime.now(),126            metadata=metadata or {}127        )128        129        self.short_term_memory[entry_id] = memory_entry130        131        # Store in interaction history132        cursor = self.db_connection.cursor()133        cursor.execute('''134            INSERT INTO interactions (id, interaction_type, content, timestamp, metadata)135            VALUES (?, ?, ?, ?, ?)136        ''', (entry_id, interaction_type, content, datetime.now(), json.dumps(metadata or {})))137        138        self.db_connection.commit()139        140        # Determine if this should be promoted to long-term memory141        importance = self._calculate_importance(content, metadata or {})142        if importance > 0.7:143            await self._promote_to_long_term(memory_entry)144    145    async def _promote_to_long_term(self, memory_entry: MemoryEntry):146        """Promote a memory entry to long-term storage"""147        # Generate embedding148        embedding = self._generate_simple_embedding(memory_entry.content)149        memory_entry.embedding = embedding150        memory_entry.memory_type = MemoryType.LONG_TERM151        152        # Store in simplified vector storage153        self.vector_storage[memory_entry.id] = {154            "content": memory_entry.content,155            "embedding": embedding,156            "timestamp": memory_entry.timestamp.isoformat(),157            "importance": memory_entry.importance,158            "metadata": json.dumps(memory_entry.metadata)159        }160    161    def _calculate_importance(self, content: str, metadata: Dict[str, Any]) -> float:162        """Calculate the importance score of a memory entry"""163        importance = 0.5  # Base importance164        165        # Increase importance for certain keywords166        important_keywords = ["error", "important", "remember", "critical", "key", "solution"]167        for keyword in important_keywords:168            if keyword.lower() in content.lower():169                importance += 0.1170        171        # Increase importance for longer content172        if len(content) > 500:173            importance += 0.1174        175        # Consider metadata factors176        if metadata.get("user_marked_important"):177            importance += 0.3178        179        return min(importance, 1.0)180    181    async def retrieve_memories(self, query: str, memory_types: List[MemoryType] = None, limit: int = 5) -> List[MemoryEntry]:182        """Retrieve relevant memories based on a query"""183        if memory_types is None:184            memory_types = [MemoryType.SHORT_TERM, MemoryType.LONG_TERM]185        186        results = []187        188        # Search short-term memory189        if MemoryType.SHORT_TERM in memory_types:190            for entry in self.short_term_memory.values():191                if query.lower() in entry.content.lower():192                    entry.access_count += 1193                    entry.last_accessed = datetime.now()194                    results.append(entry)195        196        # Search long-term memory using vector similarity197        if MemoryType.LONG_TERM in memory_types:198            query_embedding = self._generate_simple_embedding(query)199            200            # Simple similarity search in vector storage201            similarities = []202            for doc_id, doc_data in self.vector_storage.items():203                # Simple cosine similarity204                doc_embedding = doc_data['embedding']205                similarity = self._cosine_similarity(query_embedding, doc_embedding)206                similarities.append((doc_id, similarity, doc_data))207            208            # Sort by similarity and take top results209            similarities.sort(key=lambda x: x[1], reverse=True)210            211            for doc_id, similarity, doc_data in similarities[:limit]:212                results.append(MemoryEntry(213                    id=doc_id,214                    content=doc_data['content'],215                    memory_type=MemoryType.LONG_TERM,216                    timestamp=datetime.fromisoformat(doc_data['timestamp']),217                    importance=doc_data['importance'],218                    metadata=json.loads(doc_data['metadata'])219                ))220        221        # Sort by relevance and importance222        results.sort(key=lambda x: (x.importance, x.timestamp), reverse=True)223        return results[:limit]224    225    async def get_recent_interactions(self, limit: int = 10) -> List[Dict[str, Any]]:226        """Get recent interactions from memory"""227        cursor = self.db_connection.cursor()228        cursor.execute('''229            SELECT interaction_type, content, timestamp, metadata230            FROM interactions231            ORDER BY timestamp DESC232            LIMIT ?233        ''', (limit,))234        235        results = []236        for row in cursor.fetchall():237            results.append({238                "type": row[0],239                "content": row[1],240                "timestamp": row[2],241                "metadata": json.loads(row[3]) if row[3] else {}242            })243        244        return results245    246    def update_working_memory(self, key: str, value: Any):247        """Update working memory with current context"""248        self.working_memory[key] = {249            "value": value,250            "timestamp": datetime.now()251        }252    253    def get_working_memory(self, key: str = None) -> Any:254        """Get working memory content"""255        if key:256            entry = self.working_memory.get(key)257            return entry["value"] if entry else None258        return {k: v["value"] for k, v in self.working_memory.items()}259    260    def clear_working_memory(self):261        """Clear working memory"""262        self.working_memory.clear()263    264    async def store_episodic_memory(self, episode: str, context: Dict[str, Any] = None):265        """Store an episodic memory (sequence of events)"""266        entry_id = f"episode_{int(time.time() * 1000)}"267        embedding = self._generate_simple_embedding(episode)268        269        cursor = self.db_connection.cursor()270        cursor.execute('''271            INSERT INTO episodic_memory (id, content, timestamp, importance, metadata, embedding)272            VALUES (?, ?, ?, ?, ?, ?)273        ''', (274            entry_id,275            episode,276            datetime.now(),277            self._calculate_importance(episode, context or {}),278            json.dumps(context or {}),279            json.dumps(embedding)280        ))281        282        self.db_connection.commit()283    284    async def retrieve_episodic_memories(self, query: str, limit: int = 5) -> List[Dict[str, Any]]:285        """Retrieve episodic memories similar to the query"""286        query_embedding = self._generate_simple_embedding(query)287        288        cursor = self.db_connection.cursor()289        cursor.execute('''290            SELECT id, content, timestamp, importance, metadata, embedding291            FROM episodic_memory292            ORDER BY timestamp DESC293        ''')294        295        results = []296        for row in cursor.fetchall():297            stored_embedding = np.array(json.loads(row[5]))298            similarity = np.dot(query_embedding, stored_embedding) / (299                np.linalg.norm(query_embedding) * np.linalg.norm(stored_embedding)300            )301            302            if similarity > 0.5:  # Threshold for relevance303                results.append({304                    "id": row[0],305                    "content": row[1],306                    "timestamp": row[2],307                    "importance": row[3],308                    "metadata": json.loads(row[4]) if row[4] else {},309                    "similarity": float(similarity)310                })311        312        results.sort(key=lambda x: x["similarity"], reverse=True)313        return results[:limit]314    315    async def _cleanup_short_term_memory(self):316        """Periodic cleanup of expired short-term memories"""317        while True:318            current_time = datetime.now()319            expired_keys = []320            321            for key, entry in self.short_term_memory.items():322                if (current_time - entry.timestamp).seconds > self.short_term_ttl:323                    expired_keys.append(key)324            325            for key in expired_keys:326                del self.short_term_memory[key]327            328            await asyncio.sleep(300)  # Cleanup every 5 minutes329    330    async def get_memory_stats(self) -> Dict[str, Any]:331        """Get memory system statistics"""332        cursor = self.db_connection.cursor()333        334        # Count interactions335        cursor.execute("SELECT COUNT(*) FROM interactions")336        interaction_count = cursor.fetchone()[0]337        338        # Count episodic memories339        cursor.execute("SELECT COUNT(*) FROM episodic_memory")340        episodic_count = cursor.fetchone()[0]341        342        # Count vector storage entries343        vector_count = len(self.vector_storage)344        345        return {346            "short_term_entries": len(self.short_term_memory),347            "working_memory_entries": len(self.working_memory),348            "long_term_entries": vector_count,349            "episodic_entries": episodic_count,350            "total_interactions": interaction_count351        }352    353    def _generate_simple_embedding(self, text: str) -> List[float]:354        """Generate simple hash-based embedding for testing"""355        # Simple hash-based embedding (for testing without external dependencies)356        hash_obj = hashlib.md5(text.encode())357        hash_hex = hash_obj.hexdigest()358        359        # Convert hex to list of floats (normalized to 0-1 range)360        embedding = []361        for i in range(0, len(hash_hex), 2):362            val = int(hash_hex[i:i+2], 16) / 255.0363            embedding.append(val)364        365        # Pad to fixed size (16 dimensions)366        while len(embedding) < 16:367            embedding.append(0.0)368        369        return embedding[:16]370    371    def _cosine_similarity(self, vec1: List[float], vec2: List[float]) -> float:372        """Calculate cosine similarity between two vectors"""373        if len(vec1) != len(vec2):374            return 0.0375        376        dot_product = sum(a * b for a, b in zip(vec1, vec2))377        magnitude1 = sum(a * a for a in vec1) ** 0.5378        magnitude2 = sum(b * b for b in vec2) ** 0.5379        380        if magnitude1 == 0 or magnitude2 == 0:381            return 0.0382        383        return dot_product / (magnitude1 * magnitude2)384    385    async def shutdown(self):386        """Shutdown memory manager and close connections"""387        if self.db_connection:388            self.db_connection.close()389