Memory API Reference

from aakit import Agent, RedisBackend, create_memory, Message
import redis
import asyncio
from datetime import datetime, timedelta

# Example 1: Custom Memory Backend
class TimedMemoryBackend(RedisBackend):
    """Memory backend that tracks message timing"""
    
    async def save(self, session_id: str, messages: List[Message]):
        # Add timing metadata
        for msg in messages:
            msg.metadata["saved_at"] = datetime.now().isoformat()
        
        await super().save(session_id, messages)
    
    async def get_recent_messages(
        self,
        session_id: str,
        since: datetime
    ) -> List[Message]:
        """Get messages since a specific time"""
        all_messages = await self.load(session_id)
        recent = []
        
        for msg in all_messages:
            if "saved_at" in msg.metadata:
                msg_time = datetime.fromisoformat(msg.metadata["saved_at"])
                if msg_time > since:
                    recent.append(msg)
        
        return recent

# Example 2: Multi-Backend Memory
class MultiBackendMemory:
    """Use multiple backends for redundancy"""
    
    def __init__(self, backends: List[MemoryBackend]):
        self.backends = backends
    
    async def save(self, session_id: str, messages: List[Message]):
        # Save to all backends in parallel
        await asyncio.gather(*[
            backend.save(session_id, messages)
            for backend in self.backends
        ])
    
    async def load(self, session_id: str) -> List[Message]:
        # Load from first available backend
        for backend in self.backends:
            try:
                return await backend.load(session_id)
            except Exception:
                continue
        return []

# Example 3: Memory with Analytics
class AnalyticsMemory(MemoryBackend):
    """Track conversation analytics"""
    
    def __init__(self, backend: MemoryBackend):
        self.backend = backend
        self.stats = {}
    
    async def save(self, session_id: str, messages: List[Message]):
        await self.backend.save(session_id, messages)
        
        # Update analytics
        if session_id not in self.stats:
            self.stats[session_id] = {
                "message_count": 0,
                "total_tokens": 0,
                "topics": set()
            }
        
        for msg in messages:
            self.stats[session_id]["message_count"] += 1
            if "tokens" in msg.metadata:
                self.stats[session_id]["total_tokens"] += msg.metadata["tokens"]
    
    def get_analytics(self, session_id: str = None):
        if session_id:
            return self.stats.get(session_id, {})
        return self.stats

# Example 4: Complete Implementation
async def main():
    # Setup Redis backend with custom configuration
    redis_client = redis.Redis(
        host='localhost',
        port=6379,
        decode_responses=True,
        connection_pool=redis.ConnectionPool(
            max_connections=50,
            health_check_interval=30
        )
    )
    
    memory = RedisBackend(
        client=redis_client,
        prefix="agent:prod:",
        ttl=86400,  # 24 hours
        max_messages=1000,
        compression=True
    )
    
    # Create agent with memory
    agent = Agent(
        name="assistant",
        instruction="You are a helpful assistant",
        model="gpt-4",
        memory=memory
    )
    
    # Chat with session management
    session_id = "user_123"
    
    # First conversation
    response1 = await agent.chat(
        "My name is Alice and I work at TechCorp",
        session_id=session_id
    )
    
    response2 = await agent.chat(
        "What company do I work at?",
        session_id=session_id
    )
    print(response2)  # Should remember TechCorp
    
    # Search conversation history
    search_results = await memory.search(
        "TechCorp",
        session_id=session_id
    )
    
    # Get conversation statistics
    history = await memory.load(session_id)
    print(f"Total messages: {len(history)}")
    
    # Export conversation
    export_data = {
        "session_id": session_id,
        "messages": [msg.to_dict() for msg in history],
        "exported_at": datetime.now().isoformat()
    }
    
    # Clear old sessions
    await memory.clear_old_sessions(
        older_than=datetime.now() - timedelta(days=30)
    )

# Example 5: Memory Middleware
class MemoryMiddleware:
    """Add custom behavior to memory operations"""
    
    def __init__(self, memory: MemoryBackend):
        self.memory = memory
    
    async def before_save(self, session_id: str, messages: List[Message]):
        # Sanitize PII
        for msg in messages:
            msg.content = self.sanitize_pii(msg.content)
    
    async def after_load(self, session_id: str, messages: List[Message]):
        # Decrypt if needed
        for msg in messages:
            if msg.metadata.get("encrypted"):
                msg.content = self.decrypt(msg.content)
    
    def sanitize_pii(self, text: str) -> str:
        # Remove personal information
        import re
        # Remove email addresses
        text = re.sub(r'\S+@\S+', '[EMAIL]', text)
        # Remove phone numbers
        text = re.sub(r'\b\d{3}[-.]?\d{3}[-.]?\d{4}\b', '[PHONE]', text)
        return text

if __name__ == "__main__":
    asyncio.run(main())