agent.py

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