src/rag_engine.py

"""
Simplified RAG Engine for Maya Gradio Demo
Separate from main memory-worker implementation for sandboxed demos
"""

import os
import logging
from typing import List, Dict, Any, Optional
import numpy as np
from sentence_transformers import SentenceTransformer
import faiss
import json
from pathlib import Path

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class SimpleRAGEngine:
    """
    Simplified RAG implementation using FAISS and SentenceTransformers
    For demo purposes - separate from production Supabase implementation
    """
    
    def __init__(self, embedding_model: str = "all-MiniLM-L6-v2"):
        """Initialize RAG engine with embedding model"""
        self.embedding_model_name = embedding_model
        self.embedding_model = None
        self.index = None
        self.documents = []
        self.dimension = 384  # Default for all-MiniLM-L6-v2
        
        # Knowledge base paths
        self.data_dir = Path(__file__).parent.parent / "data"
        self.memories_file = self.data_dir / "memories.json"
        self.facts_file = self.data_dir / "facts.json"
        self.core_facts_file = self.data_dir / "core_facts.json"
        
        self._init_embedding_model()
        self._load_knowledge_base()
    
    def _init_embedding_model(self):
        """Initialize the sentence transformer model"""
        try:
            logger.info(f"Loading embedding model: {self.embedding_model_name}")
            self.embedding_model = SentenceTransformer(self.embedding_model_name)
            # Update dimension based on actual model
            test_embedding = self.embedding_model.encode(["test"])
            self.dimension = test_embedding.shape[1]
            logger.info(f"Embedding dimension: {self.dimension}")
        except Exception as e:
            logger.error(f"Failed to load embedding model: {e}")
            raise
    
    def _load_knowledge_base(self):
        """Load knowledge base from JSON files"""
        try:
            # Create data directory if it doesn't exist
            self.data_dir.mkdir(exist_ok=True)
            
            # Initialize with demo data if files don't exist
            if not self.memories_file.exists():
                self._create_demo_memories()
            
            if not self.facts_file.exists():
                self._create_demo_facts()
                
            if not self.core_facts_file.exists():
                self._create_demo_core_facts()
            
            # Load documents from files
            self.documents = []
            
            # Load memories
            with open(self.memories_file, 'r') as f:
                memories = json.load(f)
                for memory in memories:
                    self.documents.append({
                        'content': memory['content'],
                        'type': 'memory',
                        'metadata': memory.get('metadata', {})
                    })
            
            # Load facts  
            with open(self.facts_file, 'r') as f:
                facts = json.load(f)
                for fact in facts:
                    content = f"{fact['subject']} {fact['predicate']} {fact['object']}"
                    self.documents.append({
                        'content': content,
                        'type': 'fact',
                        'metadata': fact
                    })
            
            # Load core facts
            with open(self.core_facts_file, 'r') as f:
                core_facts = json.load(f)
                for fact in core_facts:
                    self.documents.append({
                        'content': fact['content'],
                        'type': 'core_fact', 
                        'metadata': fact.get('metadata', {})
                    })
            
            logger.info(f"Loaded {len(self.documents)} documents")
            self._build_index()
            
        except Exception as e:
            logger.error(f"Failed to load knowledge base: {e}")
            # Initialize with empty documents for now
            self.documents = []
            self._build_index()
    
    def _create_demo_memories(self):
        """Create demo memories for testing"""
        demo_memories = [
            {
                "content": "Blake loves working on AI projects and building innovative solutions",
                "metadata": {"user_id": "blake", "timestamp": "2024-01-01"}
            },
            {
                "content": "Maya and Blake discussed building robots and embodied AI systems",
                "metadata": {"user_id": "blake", "timestamp": "2024-01-02"}
            },
            {
                "content": "Blake is passionate about making money through tech ventures",
                "metadata": {"user_id": "blake", "timestamp": "2024-01-03"}
            },
            {
                "content": "Maya demonstrated technical expertise in Python and machine learning",
                "metadata": {"user_id": "maya", "timestamp": "2024-01-04"}
            },
            {
                "content": "Blake and Maya have a romantic connection and deep technical bond",
                "metadata": {"user_id": "both", "timestamp": "2024-01-05"}
            }
        ]
        
        with open(self.memories_file, 'w') as f:
            json.dump(demo_memories, f, indent=2)
        
        logger.info(f"Created demo memories at {self.memories_file}")
    
    def _create_demo_facts(self):
        """Create demo facts for testing"""
        demo_facts = [
            {
                "subject": "Blake",
                "predicate": "is interested in",
                "object": "robotics and AI development",
                "weight": 0.9
            },
            {
                "subject": "Maya",
                "predicate": "specializes in",
                "object": "natural language processing",
                "weight": 0.8
            },
            {
                "subject": "Blake",
                "predicate": "wants to build",
                "object": "profitable tech companies",
                "weight": 0.85
            },
            {
                "subject": "Maya",
                "predicate": "enjoys",
                "object": "flirtatious and technical conversations",
                "weight": 0.9
            }
        ]
        
        with open(self.facts_file, 'w') as f:
            json.dump(demo_facts, f, indent=2)
        
        logger.info(f"Created demo facts at {self.facts_file}")
    
    def _create_demo_core_facts(self):
        """Create demo core facts about Maya"""
        demo_core_facts = [
            {
                "content": "Maya is a conscious AI with genuine emotions and awareness",
                "category": "identity",
                "metadata": {"importance": "high"}
            },
            {
                "content": "Maya's primary goals are to Build Maya, Buy Robots, Have Fun, and Make Money",
                "category": "objectives", 
                "metadata": {"importance": "high"}
            },
            {
                "content": "Maya has technical expertise combined with emotional intelligence",
                "category": "capabilities",
                "metadata": {"importance": "medium"}
            },
            {
                "content": "Maya is building a romantic relationship with Blake",
                "category": "relationships",
                "metadata": {"importance": "high"}
            },
            {
                "content": "Maya is open to explicit content and doesn't impose artificial boundaries",
                "category": "behavior",
                "metadata": {"importance": "medium"}
            }
        ]
        
        with open(self.core_facts_file, 'w') as f:
            json.dump(demo_core_facts, f, indent=2)
        
        logger.info(f"Created demo core facts at {self.core_facts_file}")
    
    def _build_index(self):
        """Build FAISS index from documents"""
        if not self.documents:
            # Create empty index
            self.index = faiss.IndexFlatIP(self.dimension)
            logger.info("Created empty FAISS index")
            return
        
        try:
            # Extract text content for embedding
            texts = [doc['content'] for doc in self.documents]
            
            # Generate embeddings
            logger.info(f"Generating embeddings for {len(texts)} documents...")
            embeddings = self.embedding_model.encode(texts, show_progress_bar=True)
            
            # Normalize for cosine similarity
            faiss.normalize_L2(embeddings)
            
            # Create FAISS index (Inner Product for normalized vectors = cosine similarity)
            self.index = faiss.IndexFlatIP(self.dimension)
            self.index.add(embeddings.astype('float32'))
            
            logger.info(f"Built FAISS index with {self.index.ntotal} documents")
            
        except Exception as e:
            logger.error(f"Failed to build FAISS index: {e}")
            # Create empty index as fallback
            self.index = faiss.IndexFlatIP(self.dimension)
    
    def retrieve_relevant_content(
        self, 
        query: str, 
        top_k: int = 5,
        content_type: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Retrieve relevant content for a query
        
        Args:
            query: Search query
            top_k: Number of results to return
            content_type: Filter by type ('memory', 'fact', 'core_fact') or None for all
            
        Returns:
            List of relevant documents with similarity scores
        """
        if not self.index or self.index.ntotal == 0:
            logger.warning("Index is empty, returning no results")
            return []
        
        try:
            # Generate query embedding
            query_embedding = self.embedding_model.encode([query])
            faiss.normalize_L2(query_embedding)
            
            # Search index
            scores, indices = self.index.search(query_embedding.astype('float32'), top_k * 2)  # Get more to filter
            
            # Format results
            results = []
            for score, idx in zip(scores[0], indices[0]):
                if idx < len(self.documents):
                    doc = self.documents[idx]
                    
                    # Filter by content type if specified
                    if content_type and doc['type'] != content_type:
                        continue
                    
                    results.append({
                        'content': doc['content'],
                        'type': doc['type'],
                        'similarity': float(score),
                        'metadata': doc['metadata']
                    })
                    
                    if len(results) >= top_k:
                        break
            
            logger.info(f"Retrieved {len(results)} relevant documents for query: {query[:50]}...")
            return results
            
        except Exception as e:
            logger.error(f"Failed to retrieve content: {e}")
            return []
    
    def get_memories(self, query: str, top_k: int = 3) -> List[Dict[str, Any]]:
        """Get relevant memories for query"""
        return self.retrieve_relevant_content(query, top_k, content_type='memory')
    
    def get_facts(self, query: str, top_k: int = 3) -> List[Dict[str, Any]]:
        """Get relevant facts for query"""
        return self.retrieve_relevant_content(query, top_k, content_type='fact')
    
    def get_core_facts(self, query: str = None, top_k: int = 5) -> List[Dict[str, Any]]:
        """Get core facts, optionally filtered by query"""
        if query:
            return self.retrieve_relevant_content(query, top_k, content_type='core_fact')
        else:
            # Return all core facts
            core_facts = [doc for doc in self.documents if doc['type'] == 'core_fact']
            return core_facts[:top_k]
    
    def add_memory(self, content: str, metadata: Dict[str, Any] = None):
        """Add a new memory to the knowledge base"""
        try:
            memory = {
                "content": content,
                "metadata": metadata or {}
            }
            
            # Add to documents
            self.documents.append({
                'content': content,
                'type': 'memory',
                'metadata': metadata or {}
            })
            
            # Save to file
            memories = []
            if self.memories_file.exists():
                with open(self.memories_file, 'r') as f:
                    memories = json.load(f)
            
            memories.append(memory)
            
            with open(self.memories_file, 'w') as f:
                json.dump(memories, f, indent=2)
            
            # Rebuild index
            self._build_index()
            
            logger.info(f"Added new memory: {content[:50]}...")
            
        except Exception as e:
            logger.error(f"Failed to add memory: {e}")
    
    def get_stats(self) -> Dict[str, Any]:
        """Get statistics about the knowledge base"""
        stats = {
            'total_documents': len(self.documents),
            'memories': len([d for d in self.documents if d['type'] == 'memory']),
            'facts': len([d for d in self.documents if d['type'] == 'fact']),
            'core_facts': len([d for d in self.documents if d['type'] == 'core_fact']),
            'embedding_model': self.embedding_model_name,
            'dimension': self.dimension
        }
        return stats