ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter3/memobase/agent.py

@@ -0,0 +1,580 @@
+"""
+Memobase Agent Implementation with Kimi K3 Model
+Advanced memory management for LOCOMO benchmark
+"""
+
+import json
+import logging
+import time
+import hashlib
+import pickle
+from typing import List, Dict, Any, Optional, Tuple
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from collections import defaultdict
+from openai import OpenAI
+
+from config import (
+    KIMI_API_KEY, KIMI_BASE_URL, KIMI_MODEL,
+    MODEL_TEMPERATURE, MODEL_MAX_TOKENS, MODEL_TOP_P,
+    MEMOBASE_CONFIG, MEMORY_DB_PATH, AGENT_CONFIG,
+    MAX_MEMORY_ENTRIES, MEMORY_COMPRESSION_THRESHOLD,
+    LOG_LEVEL, LOG_FORMAT
+)
+
+
+def _reasoning_safe_temperature(model, requested=1.0):
+    """Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
+    Return 1 for those; otherwise the requested value so non-reasoning
+    providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
+    m = str(model or "").lower().replace("/", "-")
+    return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
+
+
+# Configure logging
+logging.basicConfig(level=LOG_LEVEL, format=LOG_FORMAT)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Memory:
+    """Represents a single memory entry"""
+    id: str
+    type: str  # episodic, semantic, procedural, working
+    content: Any
+    embedding: Optional[List[float]] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    created_at: datetime = field(default_factory=datetime.now)
+    accessed_at: datetime = field(default_factory=datetime.now)
+    access_count: int = 0
+    importance_score: float = 1.0
+    decay_rate: float = 0.1
+    
+    def __post_init__(self):
+        if not self.id:
+            # Generate unique ID based on content
+            content_str = json.dumps(self.content, sort_keys=True)
+            self.id = hashlib.md5(content_str.encode()).hexdigest()[:12]
+    
+    def access(self):
+        """Update access statistics"""
+        self.accessed_at = datetime.now()
+        self.access_count += 1
+        # Increase importance with access
+        self.importance_score = min(10.0, self.importance_score * 1.1)
+    
+    def decay(self):
+        """Apply time-based decay to importance"""
+        time_since_access = (datetime.now() - self.accessed_at).total_seconds() / 3600
+        self.importance_score *= (1 - self.decay_rate * min(1, time_since_access / 24))
+        self.importance_score = max(0.1, self.importance_score)
+
+
+@dataclass
+class MemoryCluster:
+    """Represents a cluster of related memories"""
+    id: str
+    memories: List[Memory]
+    summary: Optional[str] = None
+    centroid_embedding: Optional[List[float]] = None
+    created_at: datetime = field(default_factory=datetime.now)
+    
+    def add_memory(self, memory: Memory):
+        """Add a memory to the cluster"""
+        self.memories.append(memory)
+        # TODO: Update centroid embedding
+    
+    def compress(self) -> str:
+        """Compress cluster into a summary"""
+        if not self.summary:
+            # Create summary from memories
+            contents = [m.content for m in self.memories]
+            self.summary = f"Cluster of {len(self.memories)} related memories: {contents[:3]}..."
+        return self.summary
+
+
+class MemoryStore:
+    """Manages different types of memories with persistence"""
+    
+    def __init__(self, db_path: Path = MEMORY_DB_PATH):
+        self.db_path = db_path
+        self.memories: Dict[str, List[Memory]] = defaultdict(list)
+        self.clusters: List[MemoryCluster] = []
+        self.embeddings_cache: Dict[str, List[float]] = {}
+        self._load_memories()
+        logger.info(f"Initialized MemoryStore at {db_path}")
+    
+    def _load_memories(self):
+        """Load memories from persistent storage"""
+        memory_file = self.db_path / "memories.pkl"
+        if memory_file.exists():
+            try:
+                with open(memory_file, 'rb') as f:
+                    data = pickle.load(f)
+                    self.memories = data.get('memories', defaultdict(list))
+                    self.clusters = data.get('clusters', [])
+                    logger.info(f"Loaded {sum(len(m) for m in self.memories.values())} memories")
+            except Exception as e:
+                logger.error(f"Failed to load memories: {e}")
+    
+    def _save_memories(self):
+        """Save memories to persistent storage"""
+        memory_file = self.db_path / "memories.pkl"
+        try:
+            with open(memory_file, 'wb') as f:
+                pickle.dump({
+                    'memories': dict(self.memories),
+                    'clusters': self.clusters
+                }, f)
+            logger.debug("Saved memories to disk")
+        except Exception as e:
+            logger.error(f"Failed to save memories: {e}")
+    
+    def add_memory(self, memory_type: str, content: Any, 
+                   metadata: Optional[Dict] = None,
+                   importance: float = 1.0) -> Memory:
+        """Add a new memory"""
+        memory = Memory(
+            id="",  # Will be auto-generated
+            type=memory_type,
+            content=content,
+            metadata=metadata or {},
+            importance_score=importance
+        )
+        
+        self.memories[memory_type].append(memory)
+        
+        # Apply compression if needed
+        if len(self.memories[memory_type]) > MEMORY_COMPRESSION_THRESHOLD:
+            self._compress_memories(memory_type)
+        
+        self._save_memories()
+        logger.debug(f"Added {memory_type} memory: {memory.id}")
+        return memory
+    
+    def get_memories(self, memory_type: Optional[str] = None,
+                     limit: int = 10,
+                     min_importance: float = 0.5) -> List[Memory]:
+        """Retrieve memories, optionally filtered by type and importance"""
+        if memory_type:
+            memories = self.memories.get(memory_type, [])
+        else:
+            memories = [m for mlist in self.memories.values() for m in mlist]
+        
+        # Filter by importance and sort by relevance
+        memories = [m for m in memories if m.importance_score >= min_importance]
+        memories.sort(key=lambda m: (m.importance_score, m.access_count), reverse=True)
+        
+        # Update access stats
+        for memory in memories[:limit]:
+            memory.access()
+        
+        return memories[:limit]
+    
+    def search_memories(self, query: str, limit: int = 5) -> List[Memory]:
+        """Search memories by content similarity"""
+        results = []
+        query_lower = query.lower()
+        
+        for memory_list in self.memories.values():
+            for memory in memory_list:
+                # Simple text matching (would use embeddings in production)
+                content_str = str(memory.content).lower()
+                if query_lower in content_str:
+                    score = content_str.count(query_lower) * memory.importance_score
+                    results.append((score, memory))
+        
+        results.sort(key=lambda x: x[0], reverse=True)
+        
+        # Update access stats
+        for _, memory in results[:limit]:
+            memory.access()
+        
+        return [m for _, m in results[:limit]]
+    
+    def _compress_memories(self, memory_type: str):
+        """Compress old memories to save space"""
+        memories = self.memories[memory_type]
+        if len(memories) <= MEMORY_COMPRESSION_THRESHOLD:
+            return
+        
+        # Sort by importance and recency
+        memories.sort(key=lambda m: (m.importance_score, m.accessed_at.timestamp()))
+        
+        # Keep top memories, compress others
+        to_keep = memories[-MAX_MEMORY_ENTRIES//2:]
+        to_compress = memories[:-MAX_MEMORY_ENTRIES//2]
+        
+        if to_compress:
+            # Create a cluster from compressed memories
+            cluster = MemoryCluster(
+                id=hashlib.md5(f"{memory_type}_{datetime.now()}".encode()).hexdigest()[:12],
+                memories=to_compress
+            )
+            cluster.compress()
+            self.clusters.append(cluster)
+            
+            # Replace with compressed version
+            compressed_memory = Memory(
+                id=cluster.id,
+                type=memory_type,
+                content=cluster.summary,
+                metadata={"cluster_id": cluster.id, "compressed_count": len(to_compress)},
+                importance_score=sum(m.importance_score for m in to_compress) / len(to_compress)
+            )
+            
+            self.memories[memory_type] = [compressed_memory] + to_keep
+            logger.info(f"Compressed {len(to_compress)} {memory_type} memories into cluster {cluster.id}")
+    
+    def consolidate_memories(self):
+        """Consolidate and reorganize memories for efficiency"""
+        for memory_type in self.memories:
+            memories = self.memories[memory_type]
+            
+            # Apply decay to all memories
+            for memory in memories:
+                memory.decay()
+            
+            # Remove very low importance memories
+            self.memories[memory_type] = [
+                m for m in memories if m.importance_score > 0.1
+            ]
+        
+        self._save_memories()
+        logger.info("Consolidated memories")
+    
+    def clear_working_memory(self):
+        """Clear working memory (short-term)"""
+        self.memories['working'] = []
+        logger.debug("Cleared working memory")
+
+
+class MemobaseAgent:
+    """
+    Advanced agent with Memobase memory management for LOCOMO benchmark
+    """
+    
+    def __init__(self, api_key: str = KIMI_API_KEY):
+        """Initialize the Memobase agent"""
+        self.client = OpenAI(
+            api_key=api_key,
+            base_url=KIMI_BASE_URL
+        )
+        self.model = KIMI_MODEL
+        self.memory_store = MemoryStore()
+        self.conversation_history = []
+        self.current_task = None
+        self.task_context = {}
+        
+        # Initialize system prompt
+        self._init_system_prompt()
+        logger.info(f"Initialized MemobaseAgent with model {self.model}")
+    
+    def _init_system_prompt(self):
+        """Initialize the system prompt with memory capabilities"""
+        self.system_prompt = """You are an advanced AI agent with sophisticated memory management capabilities.
+
+You have access to multiple types of memory:
+1. **Episodic Memory**: Specific experiences and events from tasks
+2. **Semantic Memory**: General knowledge and facts
+3. **Procedural Memory**: Learned procedures and problem-solving patterns
+4. **Working Memory**: Current task context and temporary information
+
+Memory Management Guidelines:
+- Store important information for future reference
+- Retrieve relevant memories when solving new problems
+- Learn from past experiences to improve performance
+- Compress and consolidate memories to maintain efficiency
+- Use procedural memories to apply learned strategies
+
+Your goal is to complete tasks efficiently while learning and adapting from experience.
+When you encounter similar problems, use your memories to solve them more effectively.
+
+Always think step-by-step and use your memory system strategically."""
+    
+    def _store_interaction(self, role: str, content: str, memory_type: str = "episodic"):
+        """Store an interaction in memory"""
+        self.memory_store.add_memory(
+            memory_type=memory_type,
+            content={
+                "role": role,
+                "content": content,
+                "task": self.current_task,
+                "timestamp": datetime.now().isoformat()
+            },
+            metadata={
+                "task_id": self.current_task,
+                "turn": len(self.conversation_history)
+            }
+        )
+    
+    def _retrieve_relevant_memories(self, query: str, limit: int = 5) -> List[Memory]:
+        """Retrieve memories relevant to current query"""
+        # Search across all memory types
+        relevant_memories = []
+        
+        # Get recent episodic memories
+        episodic = self.memory_store.get_memories("episodic", limit=limit//2)
+        relevant_memories.extend(episodic)
+        
+        # Search for similar content
+        searched = self.memory_store.search_memories(query, limit=limit//2)
+        relevant_memories.extend(searched)
+        
+        # Get procedural memories if task-related
+        if "solve" in query.lower() or "how" in query.lower():
+            procedural = self.memory_store.get_memories("procedural", limit=2)
+            relevant_memories.extend(procedural)
+        
+        # Remove duplicates
+        seen = set()
+        unique_memories = []
+        for memory in relevant_memories:
+            if memory.id not in seen:
+                seen.add(memory.id)
+                unique_memories.append(memory)
+        
+        return unique_memories[:limit]
+    
+    def _format_memories_for_context(self, memories: List[Memory]) -> str:
+        """Format memories for inclusion in context"""
+        if not memories:
+            return ""
+        
+        formatted = "\n=== Relevant Memories ===\n"
+        for memory in memories:
+            formatted += f"[{memory.type.upper()}] (importance: {memory.importance_score:.2f})\n"
+            if isinstance(memory.content, dict):
+                formatted += json.dumps(memory.content, indent=2)
+            else:
+                formatted += str(memory.content)
+            formatted += "\n---\n"
+        
+        return formatted
+    
+    def _learn_from_outcome(self, task: str, approach: str, outcome: str, success: bool):
+        """Learn from task outcomes and store procedural knowledge"""
+        # Store the learning as procedural memory
+        self.memory_store.add_memory(
+            memory_type="procedural",
+            content={
+                "task_pattern": task,
+                "approach": approach,
+                "outcome": outcome,
+                "success": success,
+                "learned_at": datetime.now().isoformat()
+            },
+            importance=2.0 if success else 1.0,
+            metadata={"task_id": self.current_task}
+        )
+        
+        if success:
+            logger.info(f"Learned successful approach for task type: {task}")
+        else:
+            logger.info(f"Learned from failure in task type: {task}")
+    
+    def process_message(self, message: str, task_id: Optional[str] = None) -> str:
+        """
+        Process a message with memory-aware reasoning
+        
+        Args:
+            message: User message to process
+            task_id: Optional task identifier for context
+            
+        Returns:
+            Agent's response
+        """
+        self.current_task = task_id or f"task_{int(time.time())}"
+        
+        # Store the query in working memory
+        self.memory_store.add_memory(
+            memory_type="working",
+            content=message,
+            metadata={"task_id": self.current_task}
+        )
+        
+        # Retrieve relevant memories
+        relevant_memories = self._retrieve_relevant_memories(message)
+        memory_context = self._format_memories_for_context(relevant_memories)
+        
+        # Build messages with memory context
+        messages = [
+            {"role": "system", "content": self.system_prompt}
+        ]
+        
+        if memory_context:
+            messages.append({
+                "role": "system",
+                "content": memory_context
+            })
+        
+        # Add conversation history
+        messages.extend(self.conversation_history)
+        messages.append({"role": "user", "content": message})
+        
+        try:
+            # Call Kimi K3 model
+            response = self.client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                temperature=_reasoning_safe_temperature(self.model, MODEL_TEMPERATURE),
+                max_tokens=MODEL_MAX_TOKENS,
+                top_p=MODEL_TOP_P
+            )
+            
+            assistant_response = response.choices[0].message.content
+            
+            # Store the interaction in episodic memory
+            self._store_interaction("user", message)
+            self._store_interaction("assistant", assistant_response)
+            
+            # Update conversation history
+            self.conversation_history.append({"role": "user", "content": message})
+            self.conversation_history.append({"role": "assistant", "content": assistant_response})
+            
+            # Keep conversation history manageable
+            if len(self.conversation_history) > 20:
+                # Move old conversations to episodic memory and compress
+                old_convs = self.conversation_history[:10]
+                for conv in old_convs:
+                    self.memory_store.add_memory(
+                        memory_type="episodic",
+                        content=conv,
+                        importance=0.5
+                    )
+                self.conversation_history = self.conversation_history[10:]
+            
+            return assistant_response
+            
+        except Exception as e:
+            logger.error(f"Error processing message: {e}")
+            return f"Error: {str(e)}"
+    
+    def execute_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Execute a LOCOMO benchmark task
+        
+        Args:
+            task: Task dictionary with 'id', 'type', 'query', and optional 'context'
+            
+        Returns:
+            Result dictionary with 'response', 'memories_used', 'execution_time'
+        """
+        start_time = time.time()
+        task_id = task.get('id', f"task_{int(time.time())}")
+        task_type = task.get('type', 'unknown')
+        query = task['query']
+        context = task.get('context', '')
+        
+        self.current_task = task_id
+        
+        # Check for similar past tasks in procedural memory
+        similar_tasks = self.memory_store.search_memories(f"{task_type} {query[:50]}", limit=3)
+        
+        # Build enhanced query with context
+        enhanced_query = query
+        if context:
+            enhanced_query = f"Context: {context}\n\nTask: {query}"
+        
+        # Process the task
+        response = self.process_message(enhanced_query, task_id)
+        
+        # Extract approach and outcome for learning
+        approach = f"Used {len(similar_tasks)} similar memories"
+        outcome = response[:100]  # First 100 chars as outcome summary
+        
+        # Learn from this task
+        self._learn_from_outcome(
+            task=task_type,
+            approach=approach,
+            outcome=outcome,
+            success=True  # Would be determined by evaluation
+        )
+        
+        execution_time = time.time() - start_time
+        
+        return {
+            "task_id": task_id,
+            "response": response,
+            "memories_used": len(similar_tasks),
+            "execution_time": execution_time,
+            "memory_stats": {
+                "episodic": len(self.memory_store.memories.get('episodic', [])),
+                "semantic": len(self.memory_store.memories.get('semantic', [])),
+                "procedural": len(self.memory_store.memories.get('procedural', [])),
+                "working": len(self.memory_store.memories.get('working', []))
+            }
+        }
+    
+    def consolidate_and_learn(self):
+        """Consolidate memories and extract learnings"""
+        logger.info("Starting memory consolidation...")
+        
+        # Consolidate memories
+        self.memory_store.consolidate_memories()
+        
+        # Extract patterns from episodic memories
+        episodic_memories = self.memory_store.get_memories('episodic', limit=50)
+        
+        # Group by task type and extract patterns
+        task_patterns = defaultdict(list)
+        for memory in episodic_memories:
+            if isinstance(memory.content, dict):
+                task = memory.content.get('task', 'unknown')
+                task_patterns[task].append(memory)
+        
+        # Create procedural memories from patterns
+        for task_type, memories in task_patterns.items():
+            if len(memories) >= 3:  # Need multiple examples to learn
+                # Extract common approach
+                pattern = {
+                    "task_type": task_type,
+                    "successful_approaches": [],
+                    "common_challenges": [],
+                    "learned_from": len(memories)
+                }
+                
+                self.memory_store.add_memory(
+                    memory_type="procedural",
+                    content=pattern,
+                    importance=2.0,
+                    metadata={"consolidation_run": datetime.now().isoformat()}
+                )
+        
+        # Clear working memory
+        self.memory_store.clear_working_memory()
+        
+        logger.info("Memory consolidation complete")
+    
+    def reset(self, keep_memories: bool = True):
+        """Reset the agent state"""
+        self.conversation_history = []
+        self.current_task = None
+        self.task_context = {}
+        
+        if not keep_memories:
+            self.memory_store = MemoryStore()
+        else:
+            # Only clear working memory
+            self.memory_store.clear_working_memory()
+        
+        logger.info(f"Agent reset (memories kept: {keep_memories})")
+    
+    def get_performance_metrics(self) -> Dict[str, Any]:
+        """Get agent performance metrics"""
+        memory_stats = {
+            memory_type: len(memories)
+            for memory_type, memories in self.memory_store.memories.items()
+        }
+        
+        total_memories = sum(memory_stats.values())
+        cluster_count = len(self.memory_store.clusters)
+        
+        return {
+            "total_memories": total_memories,
+            "memory_distribution": memory_stats,
+            "clusters_created": cluster_count,
+            "conversation_length": len(self.conversation_history),
+            "current_task": self.current_task
+        }