long_term_memory.py

"""
Long-term memory implementation using ChromaDB for persistent storage
Following specifications from rules.mdc
"""

import os
import uuid
import json
import logging
from datetime import datetime, timedelta
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass, asdict
from enum import Enum

import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
from openai import OpenAI


class MemoryType(Enum):
    """Types of memories as defined in specs"""
    SEMANTIC = "semantic"      # Facts and preferences  
    EPISODIC = "episodic"      # What happened
    PROCEDURAL = "procedural"  # Routines and how-to


class SensitivityLevel(Enum):
    """Sensitivity levels for memories"""
    LOW = "low"
    MEDIUM = "medium" 
    HIGH = "high"


@dataclass
class MemoryEntry:
    """Memory entry schema from rules.mdc"""
    id: str
    user_id: str
    type: MemoryType
    text: str
    embedding_id: str
    domain: str
    sensitivity: SensitivityLevel
    source: str
    created_at: str
    ttl_days: int
    consent_flag: bool
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for storage"""
        data = asdict(self)
        data['type'] = self.type.value
        data['sensitivity'] = self.sensitivity.value
        return data
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'MemoryEntry':
        """Create from dictionary"""
        data['type'] = MemoryType(data['type'])
        data['sensitivity'] = SensitivityLevel(data['sensitivity'])
        return cls(**data)
    
    def is_expired(self) -> bool:
        """Check if memory has expired based on TTL"""
        created = datetime.fromisoformat(self.created_at)
        expiry = created + timedelta(days=self.ttl_days)
        return datetime.now() > expiry


class MemoryWriter:
    """Extracts salient facts from conversations using LLM"""
    
    def __init__(self, openai_client: OpenAI):
        self.client = openai_client
        self.extraction_prompt = """You are a memory extraction system. Analyze the conversation turn and extract important, novel information that should be remembered long-term.

EXTRACTION CRITERIA:
- Extract user preferences, interests, personal facts
- Extract significant events or experiences mentioned
- Extract useful procedures or how-to information  
- IGNORE trivialities, small talk, or temporary information
- ONLY extract if information is novel and worth remembering

DOMAINS: preferences, interests, work, health, finance, personal, general

For each significant piece of information, classify as:
- semantic: Facts, preferences, traits ("User likes coffee", "User works in tech")
- episodic: Events, experiences ("User went to Tokyo last month")  
- procedural: Routines, processes ("User's morning routine includes yoga")

IMPORTANT: Return ONLY a valid JSON array. No explanations, no code blocks, no other text.
If nothing significant to remember, return: []

Format:
[
  {{
    "text": "Brief, clear description of the fact",
    "type": "semantic",
    "domain": "preferences",
    "sensitivity": "low"
  }}
]

Conversation turn:
User: {user_input}
Assistant: {assistant_response}

JSON array:"""

    def extract_memories(self, user_input: str, assistant_response: str, user_id: str, source: str) -> List[MemoryEntry]:
        """Extract salient memories from a conversation turn"""
        try:
            # Generate extraction prompt
            prompt = self.extraction_prompt.format(
                user_input=user_input,
                assistant_response=assistant_response
            )
            
            # Call LLM for extraction
            response = self.client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": prompt}],
                max_tokens=500,
                temperature=0.1
            )
            
            # Parse response
            content = response.choices[0].message.content.strip()
            if not content or content.lower().startswith('[]'):
                return []
            
            # Parse JSON response
            try:
                # Clean up the response first
                content = content.strip()
                if content.startswith('```json'):
                    content = content[7:]
                if content.endswith('```'):
                    content = content[:-3]
                content = content.strip()
                
                extracted_data = json.loads(content)
            except json.JSONDecodeError as e:
                # Try to extract JSON from response if wrapped in text
                import re
                json_match = re.search(r'\[.*?\]', content, re.DOTALL)
                if json_match:
                    try:
                        extracted_data = json.loads(json_match.group(0))
                    except json.JSONDecodeError:
                        logging.warning(f"Failed to parse extracted JSON: {json_match.group(0)}")
                        return []
                else:
                    logging.warning(f"Failed to parse memory extraction: {content[:200]}...")
                    return []
            
            # Convert to MemoryEntry objects
            memories = []
            for item in extracted_data:
                # Set TTL based on type
                ttl_days = self._get_ttl_for_type(MemoryType(item['type']))
                
                memory = MemoryEntry(
                    id=str(uuid.uuid4()),
                    user_id=user_id,
                    type=MemoryType(item['type']),
                    text=item['text'],
                    embedding_id="",  # Will be set by LongTermMemory
                    domain=item['domain'],
                    sensitivity=SensitivityLevel(item['sensitivity']),
                    source=source,
                    created_at=datetime.now().isoformat(),
                    ttl_days=ttl_days,
                    consent_flag=True  # For now, assume consent
                )
                memories.append(memory)
            
            return memories
            
        except Exception as e:
            logging.error(f"Memory extraction failed: {e}")
            return []
    
    def _get_ttl_for_type(self, memory_type: MemoryType) -> int:
        """Get TTL days based on memory type (from specs)"""
        ttl_map = {
            MemoryType.SEMANTIC: 365,    # 1 year
            MemoryType.EPISODIC: 180,    # 6 months  
            MemoryType.PROCEDURAL: 365   # 1 year
        }
        return ttl_map.get(memory_type, 180)


class LongTermMemory:
    """ChromaDB-based long-term memory system"""
    
    def __init__(self, persist_directory: str = "./chroma_db", user_id: str = "default_user"):
        self.user_id = user_id
        self.persist_directory = persist_directory
        
        # Initialize ChromaDB client
        self.client = chromadb.PersistentClient(
            path=persist_directory,
            settings=Settings(
                anonymized_telemetry=False,
                allow_reset=True
            )
        )
        
        # Initialize embedding model
        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        
        # Get or create collections (namespaced per user)
        self.collection_name = f"personal_memory_{user_id}"
        try:
            self.collection = self.client.get_collection(self.collection_name)
        except:
            self.collection = self.client.create_collection(
                name=self.collection_name,
                metadata={"description": f"Personal memories for user {user_id}"}
            )
        
        logging.info(f"LongTermMemory initialized for user {user_id}")
    
    def store_memories(self, memories: List[MemoryEntry]) -> int:
        """Store memories in vector database"""
        if not memories:
            return 0
        
        stored_count = 0
        for memory in memories:
            try:
                # Check for duplicates using semantic similarity
                if self._is_duplicate(memory):
                    logging.info(f"Skipping duplicate memory: {memory.text}")
                    continue
                
                # Generate embedding
                embedding = self.embedding_model.encode(memory.text).tolist()
                
                # Set embedding ID
                memory.embedding_id = f"vec_{memory.id}"
                
                # Store in ChromaDB
                self.collection.add(
                    documents=[memory.text],
                    embeddings=[embedding],
                    metadatas=[memory.to_dict()],
                    ids=[memory.id]
                )
                
                stored_count += 1
                logging.info(f"Stored memory: {memory.text[:50]}...")
                
            except Exception as e:
                logging.error(f"Failed to store memory: {e}")
        
        return stored_count
    
    def retrieve_memories(self, query: str, k: int = 3, domain_filters: Optional[List[str]] = None) -> List[Tuple[MemoryEntry, float]]:
        """Retrieve relevant memories for a query"""
        try:
            # Generate query embedding
            query_embedding = self.embedding_model.encode(query).tolist()
            
            # Build metadata filters - simplified for ChromaDB compatibility
            where_clause = {"user_id": self.user_id}
            
            # Query ChromaDB (simplified without domain filters for now)
            results = self.collection.query(
                query_embeddings=[query_embedding],
                n_results=k * 2,  # Get more results to filter later
                where=where_clause
            )
            
            # Post-filter by domain if needed
            if domain_filters:
                filtered_docs = []
                filtered_metadatas = []
                filtered_distances = []
                
                for i, metadata in enumerate(results['metadatas'][0]):
                    if metadata.get('domain') in domain_filters:
                        filtered_docs.append(results['documents'][0][i])
                        filtered_metadatas.append(metadata)
                        if results['distances']:
                            filtered_distances.append(results['distances'][0][i])
                
                # Rebuild results structure
                results = {
                    'documents': [filtered_docs[:k]],
                    'metadatas': [filtered_metadatas[:k]],
                    'distances': [filtered_distances[:k]] if filtered_distances else None
                }
            
            if not results['documents'] or not results['documents'][0]:
                return []
            
            # Process results
            memories = []
            for i, doc in enumerate(results['documents'][0]):
                metadata = results['metadatas'][0][i]
                distance = results['distances'][0][i] if results['distances'] else [0.0]
                
                # Create MemoryEntry from metadata
                memory = MemoryEntry.from_dict(metadata)
                
                # Check if memory is expired
                if memory.is_expired():
                    continue
                
                # Calculate similarity score (1 - distance)
                similarity_score = 1.0 - distance[i] if isinstance(distance, list) else 1.0 - distance
                
                memories.append((memory, similarity_score))
            
            # Sort by similarity + recency boost (recent memories get slight boost)
            memories.sort(key=lambda x: self._calculate_relevance_score(x[0], x[1]), reverse=True)
            
            return memories
            
        except Exception as e:
            logging.error(f"Memory retrieval failed: {e}")
            return []
    
    def get_conversation_context(self, current_query: str, max_memories: int = 3) -> str:
        """Get formatted conversation context from relevant memories"""
        relevant_memories = self.retrieve_memories(current_query, k=max_memories)
        
        if not relevant_memories:
            return "No relevant memories found."
        
        context_parts = ["Relevant memories from previous conversations:"]
        for memory, score in relevant_memories:
            # Format memory with recency info
            created_date = datetime.fromisoformat(memory.created_at).strftime("%Y-%m-%d")
            context_parts.append(f"- {memory.text} (from {created_date})")
        
        return "\n".join(context_parts)
    
    def cleanup_expired_memories(self) -> int:
        """Remove expired memories"""
        try:
            # Get all memories for this user
            all_results = self.collection.get(
                where={"user_id": self.user_id}
            )
            
            expired_ids = []
            for i, metadata in enumerate(all_results['metadatas']):
                memory = MemoryEntry.from_dict(metadata)
                if memory.is_expired():
                    expired_ids.append(all_results['ids'][i])
            
            # Delete expired memories
            if expired_ids:
                self.collection.delete(ids=expired_ids)
                logging.info(f"Cleaned up {len(expired_ids)} expired memories")
            
            return len(expired_ids)
            
        except Exception as e:
            logging.error(f"Memory cleanup failed: {e}")
            return 0
    
    def get_memory_stats(self) -> Dict[str, Any]:
        """Get statistics about stored memories"""
        try:
            all_results = self.collection.get(
                where={"user_id": self.user_id}
            )
            
            total_memories = len(all_results['ids'])
            
            # Count by type and domain
            type_counts = {}
            domain_counts = {}
            
            for metadata in all_results['metadatas']:
                memory_type = metadata.get('type', 'unknown')
                domain = metadata.get('domain', 'unknown')
                
                type_counts[memory_type] = type_counts.get(memory_type, 0) + 1
                domain_counts[domain] = domain_counts.get(domain, 0) + 1
            
            return {
                'user_id': self.user_id,
                'total_memories': total_memories,
                'by_type': type_counts,
                'by_domain': domain_counts
            }
            
        except Exception as e:
            logging.error(f"Failed to get memory stats: {e}")
            return {'error': str(e)}
    
    def _is_duplicate(self, new_memory: MemoryEntry, similarity_threshold: float = 0.85) -> bool:
        """Check if memory is a duplicate using semantic similarity"""
        try:
            # Search for similar memories
            similar = self.retrieve_memories(new_memory.text, k=3, domain_filters=[new_memory.domain])
            
            for existing_memory, score in similar:
                if score > similarity_threshold:
                    return True
            
            return False
            
        except Exception as e:
            logging.warning(f"Duplicate check failed: {e}")
            return False
    
    def _calculate_relevance_score(self, memory: MemoryEntry, semantic_score: float) -> float:
        """Calculate relevance score combining semantic similarity and recency"""
        # Base semantic score
        relevance = semantic_score
        
        # Recency boost (more recent memories get slight boost)
        try:
            created = datetime.fromisoformat(memory.created_at)
            days_ago = (datetime.now() - created).days
            recency_boost = max(0, (30 - days_ago) / 30 * 0.1)  # Up to 0.1 boost for memories < 30 days
            relevance += recency_boost
        except:
            pass
        
        # Domain priority (could be expanded based on user settings)
        domain_boosts = {
            'preferences': 0.05,
            'interests': 0.03,
            'personal': 0.02
        }
        relevance += domain_boosts.get(memory.domain, 0)
        
        return relevance


# Example usage and testing
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    # Test memory system
    openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    
    # Initialize components
    memory_writer = MemoryWriter(openai_client)
    ltm = LongTermMemory(user_id="test_user")
    
    # Test memory extraction and storage
    test_conversation = {
        "user_input": "I love drinking pour-over coffee every morning, especially Ethiopian beans.",
        "assistant_response": "That's wonderful! Pour-over coffee really brings out the unique flavors of Ethiopian beans. Do you have a favorite brewing method?"
    }
    
    # Extract memories
    memories = memory_writer.extract_memories(
        test_conversation["user_input"],
        test_conversation["assistant_response"],
        "test_user",
        "test_conversation"
    )
    
    print(f"Extracted {len(memories)} memories:")
    for memory in memories:
        print(f"- {memory.text} ({memory.type.value}, {memory.domain})")
    
    # Store memories
    stored_count = ltm.store_memories(memories)
    print(f"Stored {stored_count} memories")
    
    # Test retrieval
    relevant = ltm.retrieve_memories("coffee preferences")
    print(f"\nRetrieved {len(relevant)} relevant memories:")
    for memory, score in relevant:
        print(f"- {memory.text} (score: {score:.3f})")
    
    # Get stats
    stats = ltm.get_memory_stats()
    print(f"\nMemory stats: {stats}")