kite-agent 0.1.0__py3-none-any.whl

kite/memory/session_memory.py

@@ -0,0 +1,423 @@
+"""
+Session Memory Manager
+Based on Chapter 3.2: Short-Term Memory - Managing the "Now"
+
+The Goldfish Metaphor from book:
+- Sharp about 2 minutes ago
+- Total amnesia about 5 minutes ago
+
+Strategy: Sliding window with compression to avoid exponential cost growth.
+
+Run: python session_memory.py
+"""
+
+import os
+from typing import List, Dict, Optional
+from dataclasses import dataclass, field
+from datetime import datetime
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+@dataclass
+class Message:
+    """A conversation message."""
+    role: str  # 'user' or 'assistant'
+    content: str
+    timestamp: datetime = field(default_factory=datetime.now)
+    tokens: int = 0
+
+
+@dataclass
+class SessionStats:
+    """Statistics for a session."""
+    total_messages: int = 0
+    total_tokens: int = 0
+    compressed_count: int = 0
+    window_size: int = 0
+
+
+class SessionMemory:
+    """
+    Sliding window memory manager for conversations.
+    
+    Features from Chapter 3.2:
+    - Keep last N messages (configurable window)
+    - Compress older messages for context
+    - Fixed cost per conversation
+    - No exponential growth
+    
+    Example:
+        memory = SessionMemoryManager(window_size=10)
+        
+        # Add messages
+        memory.add_user_message("Hello")
+        memory.add_assistant_message("Hi! How can I help?")
+        
+        # Get messages for LLM
+        messages = memory.get_messages()
+    """
+    
+    def __init__(
+        self,
+        llm = None,
+        window_size: int = 10,
+        compression_enabled: bool = True,
+        max_tokens_per_message: int = 500
+    ):
+        """
+        Initialize session memory.
+        
+        Args:
+            window_size: Number of recent messages to keep
+            compression_enabled: Whether to compress old messages
+            max_tokens_per_message: Soft limit for message length
+        """
+        self.llm = llm
+        self.window_size = window_size
+        self.compression_enabled = compression_enabled
+        self.max_tokens_per_message = max_tokens_per_message
+        
+        # Message storage
+        self.messages: List[Message] = []
+        
+        # Compressed history (for context)
+        self.compressed_history: Optional[str] = None
+        
+        # Statistics
+        self.stats = SessionStats(window_size=window_size)
+        
+        print(f"[OK] Session Memory initialized")
+        print(f"  Window size: {window_size} messages")
+        print(f"  Compression: {'enabled' if compression_enabled else 'disabled'}")
+    
+    def _estimate_tokens(self, text: str) -> int:
+        """
+        Estimate token count.
+        
+        Rough estimation: 1 token   4 characters in English.
+        For production, use tiktoken library.
+        """
+        return len(text) // 4
+    
+    def _compress_messages(self, messages: List[Message]) -> str:
+        """
+        Compress messages into summary.
+        
+        Uses LLM to create concise summary of conversation history.
+        This is much cheaper than keeping all messages.
+        """
+        if not messages:
+            return ""
+        
+        # Build text from messages
+        conversation = "\n".join([
+            f"{msg.role.upper()}: {msg.content}"
+            for msg in messages
+        ])
+        
+        # Ask LLM to summarize
+        prompt = f"""Summarize this conversation history in 2-3 sentences. Focus on key topics and decisions.
+
+Conversation:
+{conversation}
+
+Summary:"""
+        
+        if self.llm:
+            response = self.llm.complete(prompt, max_tokens=150, temperature=0.3)
+            summary = response
+        else:
+            # Fallback for demo/testing if no LLM provided
+            summary = "Summary unavailable (No LLM provider)"
+        
+        print(f"    Compressed {len(messages)} messages   {len(summary)} chars")
+        
+        return summary
+    
+    def add_message(self, role: str, content: str, session_id: Optional[str] = None):
+        """Add a message to memory with a specific role."""
+        if role.lower() == "user":
+            self.add_user_message(content)
+        else:
+            self.add_assistant_message(content)
+
+    def add_user_message(self, content: str):
+        """Add user message to memory."""
+        message = Message(
+            role="user",
+            content=content,
+            tokens=self._estimate_tokens(content)
+        )
+        
+        self.messages.append(message)
+        self.stats.total_messages += 1
+        self.stats.total_tokens += message.tokens
+        
+        # Apply sliding window
+        self._apply_sliding_window()
+    
+    def add_assistant_message(self, content: str):
+        """Add assistant message to memory."""
+        message = Message(
+            role="assistant",
+            content=content,
+            tokens=self._estimate_tokens(content)
+        )
+        
+        self.messages.append(message)
+        self.stats.total_messages += 1
+        self.stats.total_tokens += message.tokens
+        
+        # Apply sliding window
+        self._apply_sliding_window()
+    
+    def _apply_sliding_window(self):
+        """
+        Apply sliding window logic.
+        
+        From Chapter 3.2:
+        - Keep last N messages (sharp memory)
+        - Compress older messages (compressed context)
+        - Total amnesia for very old (to save cost)
+        """
+        if len(self.messages) <= self.window_size:
+            # Within window, no action needed
+            return
+        
+        # Messages outside window
+        old_messages = self.messages[:-self.window_size]
+        
+        # Keep only window
+        self.messages = self.messages[-self.window_size:]
+        
+        # Compress old messages if enabled
+        if self.compression_enabled and old_messages:
+            new_compression = self._compress_messages(old_messages)
+            
+            # Merge with existing compression
+            if self.compressed_history:
+                # Combine old and new compression
+                combined = f"{self.compressed_history}\n{new_compression}"
+                # Compress the compression if it gets too long
+                if len(combined) > 1000:
+                    self.compressed_history = self._compress_messages([
+                        Message(role="system", content=combined)
+                    ])
+                else:
+                    self.compressed_history = combined
+            else:
+                self.compressed_history = new_compression
+            
+            self.stats.compressed_count += len(old_messages)
+            
+            print(f"      Sliding window applied: {len(old_messages)} messages compressed")
+    
+    def get_messages(self, include_compression: bool = True) -> List[Dict]:
+        """
+        Get messages in format for LLM API.
+        
+        Args:
+            include_compression: Whether to include compressed history
+            
+        Returns:
+            List of message dictionaries
+        """
+        result = []
+        
+        # Add compressed history as system message
+        if include_compression and self.compressed_history:
+            result.append({
+                "role": "system",
+                "content": f"Previous conversation summary: {self.compressed_history}"
+            })
+        
+        # Add current window
+        for msg in self.messages:
+            result.append({
+                "role": msg.role,
+                "content": msg.content
+            })
+        
+        return result
+    
+    def get_token_count(self) -> Dict[str, int]:
+        """
+        Get token counts for cost estimation.
+        
+        This shows the cost savings from sliding window approach.
+        """
+        # Current window tokens
+        window_tokens = sum(msg.tokens for msg in self.messages)
+        
+        # Compressed history tokens
+        compression_tokens = self._estimate_tokens(self.compressed_history or "")
+        
+        # What it would be WITHOUT compression
+        naive_tokens = self.stats.total_tokens
+        
+        # Actual tokens with compression
+        actual_tokens = window_tokens + compression_tokens
+        
+        return {
+            "window_tokens": window_tokens,
+            "compression_tokens": compression_tokens,
+            "actual_tokens": actual_tokens,
+            "naive_tokens": naive_tokens,
+            "savings": naive_tokens - actual_tokens,
+            "savings_percent": (
+                ((naive_tokens - actual_tokens) / naive_tokens * 100)
+                if naive_tokens > 0 else 0
+            )
+        }
+    
+    def clear(self):
+        """Clear all memory."""
+        self.messages.clear()
+        self.compressed_history = None
+        self.stats = SessionStats(window_size=self.window_size)
+        print("[OK] Session memory cleared")
+    
+    def get_stats(self) -> SessionStats:
+        """Get session statistics."""
+        return self.stats
+
+
+# ============================================================================
+# DEMO
+# ============================================================================
+
+def demo():
+    print("=" * 70)
+    print("SESSION MEMORY MANAGER DEMO")
+    print("=" * 70)
+    print("\nBased on Chapter 3.2: The Goldfish Metaphor")
+    print("- Sharp about recent messages")
+    print("- Compressed for older context")
+    print("- Total amnesia for very old (cost savings!)")
+    print("=" * 70)
+    
+    # Initialize with small window for demo
+    memory = SessionMemoryManager(
+        window_size=6,  # Keep last 6 messages
+        compression_enabled=True
+    )
+    
+    # Simulate a long conversation
+    conversation = [
+        ("user", "Hi! I need help with my Python code."),
+        ("assistant", "Of course! I'd be happy to help. What's the issue?"),
+        ("user", "I'm getting a TypeError when I try to concatenate strings."),
+        ("assistant", "That usually happens when you try to concatenate a string with a non-string type. Can you show me the code?"),
+        ("user", "Sure: result = 'Count: ' + count"),
+        ("assistant", "I see the issue! 'count' is probably an integer. You need to convert it: result = 'Count: ' + str(count)"),
+        ("user", "That worked! Thanks!"),
+        ("assistant", "Great! Is there anything else you need help with?"),
+        ("user", "Actually yes, how do I read a CSV file?"),
+        ("assistant", "You can use the csv module or pandas. Here's a simple example with csv module..."),
+        ("user", "What about using pandas instead?"),
+        ("assistant", "With pandas it's even easier: import pandas as pd; df = pd.read_csv('file.csv')"),
+        ("user", "Perfect! One more thing - how do I handle errors?"),
+        ("assistant", "Use try-except blocks to handle exceptions gracefully..."),
+    ]
+    
+    print("\n  Simulating 14-message conversation...")
+    print(f"   Window size: {memory.window_size} messages\n")
+    
+    for i, (role, content) in enumerate(conversation, 1):
+        print(f"{i}. {role.upper()}: {content[:50]}...")
+        
+        if role == "user":
+            memory.add_user_message(content)
+        else:
+            memory.add_assistant_message(content)
+        
+        # Show window status every few messages
+        if i % 4 == 0:
+            print(f"\n   [CHART] After {i} messages:")
+            print(f"      In window: {len(memory.messages)}")
+            print(f"      Compressed: {memory.stats.compressed_count}")
+            if memory.compressed_history:
+                print(f"      Compression: {len(memory.compressed_history)} chars")
+            print()
+    
+    # Show final state
+    print("\n" + "="*70)
+    print("FINAL MEMORY STATE")
+    print("="*70)
+    
+    messages = memory.get_messages()
+    
+    if memory.compressed_history:
+        print(f"\n  Compressed History:")
+        print(f"   {memory.compressed_history}\n")
+    
+    print(f"  Current Window ({len(memory.messages)} messages):")
+    for msg in memory.messages:
+        print(f"   {msg.role.upper()}: {msg.content[:60]}...")
+    
+    # Token analysis
+    print("\n" + "="*70)
+    print("  COST ANALYSIS")
+    print("="*70)
+    
+    tokens = memory.get_token_count()
+    
+    print(f"\nToken Counts:")
+    print(f"  Window tokens: {tokens['window_tokens']:,}")
+    print(f"  Compression tokens: {tokens['compression_tokens']:,}")
+    print(f"  Actual total: {tokens['actual_tokens']:,}")
+    print()
+    print(f"Without Compression:")
+    print(f"  Naive total: {tokens['naive_tokens']:,}")
+    print()
+    print(f"  Savings:")
+    print(f"  Tokens saved: {tokens['savings']:,}")
+    print(f"  Percentage: {tokens['savings_percent']:.1f}%")
+    
+    # Cost estimate
+    cost_per_1k = 0.0001  # Rough estimate
+    actual_cost = (tokens['actual_tokens'] / 1000) * cost_per_1k
+    naive_cost = (tokens['naive_tokens'] / 1000) * cost_per_1k
+    
+    print(f"\nCost Estimate (at ${cost_per_1k:.4f} per 1K tokens):")
+    print(f"  With sliding window: ${actual_cost:.6f} per query")
+    print(f"  Without compression: ${naive_cost:.6f} per query")
+    print(f"  Savings: ${naive_cost - actual_cost:.6f} per query")
+    
+    # Scale to many queries
+    queries_per_day = 1000
+    print(f"\nFor {queries_per_day:,} queries/day (30 days):")
+    print(f"  With sliding window: ${actual_cost * queries_per_day * 30:.2f}/month")
+    print(f"  Without: ${naive_cost * queries_per_day * 30:.2f}/month")
+    print(f"    Total savings: ${(naive_cost - actual_cost) * queries_per_day * 30:.2f}/month")
+    
+    print("\n" + "="*70)
+    print("KEY INSIGHTS FROM CHAPTER 3.2")
+    print("="*70)
+    print("""
+1. GOLDFISH MEMORY STRATEGY
+   - Keep recent messages (sharp memory)
+   - Compress older context
+   - Forget very old (cost savings)
+
+2. COST CONTROL
+   - Fixed cost per conversation
+   - No exponential growth
+   - Predictable budgeting
+
+3. QUALITY MAINTAINED
+   - Important context preserved via compression
+   - Recent messages at full detail
+   - Good balance of context vs. cost
+
+4. SCALABILITY
+   - Window size adjustable
+   - Compression depth configurable
+   - Works for short or long conversations
+    """)
+
+
+if __name__ == "__main__":
+    demo()