Create utils/llamaindex_rag.py

utils/llamaindex_rag.py

@@ -0,0 +1,309 @@
+"""
+LlamaIndex RAG for Credit Card Benefits Knowledge Base
+Provides intelligent context for card recommendations
+"""
+
+import os
+import logging
+from typing import Optional, Dict, List
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+# Check if LlamaIndex is available
+LLAMAINDEX_AVAILABLE = False
+try:
+    from llama_index.core import (
+        VectorStoreIndex, 
+        SimpleDirectoryReader, 
+        Settings,
+        StorageContext,
+        load_index_from_storage
+    )
+    from llama_index.embeddings.openai import OpenAIEmbedding
+    from llama_index.llms.openai import OpenAI
+    LLAMAINDEX_AVAILABLE = True
+    logger.info("✅ LlamaIndex library imported successfully")
+except ImportError as e:
+    logger.warning(f"⚠️ LlamaIndex not installed: {e}")
+    logger.warning("Install with: pip install llama-index llama-index-embeddings-openai llama-index-llms-openai")
+except Exception as e:
+    logger.error(f"❌ Error importing LlamaIndex: {e}")
+
+
+class CardBenefitsRAG:
+    """RAG system for credit card benefits using LlamaIndex"""
+    
+    def __init__(self, data_dir: str = "data/card_benefits", persist_dir: str = ".index_storage"):
+        """
+        Initialize LlamaIndex RAG
+        
+        Args:
+            data_dir: Directory containing card benefit markdown files
+            persist_dir: Directory to persist the index
+        """
+        self.enabled = LLAMAINDEX_AVAILABLE and bool(os.getenv("OPENAI_API_KEY"))
+        
+        if not LLAMAINDEX_AVAILABLE:
+            logger.warning("⚠️ LlamaIndex library not available")
+            logger.warning("To enable RAG features:")
+            logger.warning("1. Add 'llama-index' to requirements.txt")
+            logger.warning("2. Run: pip install llama-index llama-index-embeddings-openai llama-index-llms-openai")
+            return
+        
+        if not os.getenv("OPENAI_API_KEY"):
+            logger.warning("⚠️ OPENAI_API_KEY not found in environment")
+            logger.warning("RAG features will be disabled")
+            return
+        
+        self.data_dir = Path(data_dir)
+        self.persist_dir = Path(persist_dir)
+        
+        try:
+            # Configure LlamaIndex settings
+            Settings.llm = OpenAI(
+                model="gpt-4-turbo-preview",
+                temperature=0.1,
+                api_key=os.getenv("OPENAI_API_KEY")
+            )
+            Settings.embed_model = OpenAIEmbedding(
+                model="text-embedding-3-small",
+                api_key=os.getenv("OPENAI_API_KEY")
+            )
+            Settings.chunk_size = 512
+            Settings.chunk_overlap = 50
+            
+            # Load or create index
+            if self.persist_dir.exists():
+                logger.info("📚 Loading existing LlamaIndex from storage...")
+                try:
+                    storage_context = StorageContext.from_defaults(persist_dir=str(self.persist_dir))
+                    self.index = load_index_from_storage(storage_context)
+                    logger.info("✅ Index loaded from storage")
+                except Exception as e:
+                    logger.warning(f"⚠️ Could not load index from storage: {e}")
+                    logger.info("Creating new index...")
+                    self._create_index()
+            else:
+                logger.info("📚 Creating new LlamaIndex from documents...")
+                self._create_index()
+            
+            # Create query engine
+            self.query_engine = self.index.as_query_engine(
+                similarity_top_k=3,
+                response_mode="compact"
+            )
+            
+            logger.info("✅ CardBenefitsRAG initialized successfully")
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to initialize CardBenefitsRAG: {e}")
+            import traceback
+            traceback.print_exc()
+            self.enabled = False
+    
+    def _create_index(self):
+        """Create index from documents"""
+        if not self.data_dir.exists():
+            logger.warning(f"⚠️ Data directory not found: {self.data_dir}")
+            logger.info("Creating data directory with sample file...")
+            self.data_dir.mkdir(parents=True, exist_ok=True)
+            self._create_sample_data()
+        
+        # Check if directory has any files
+        files = list(self.data_dir.glob("*.md")) + list(self.data_dir.glob("*.txt"))
+        if not files:
+            logger.warning("⚠️ No markdown or text files found in data directory")
+            logger.info("Creating sample file...")
+            self._create_sample_data()
+        
+        # Load documents
+        try:
+            documents = SimpleDirectoryReader(
+                str(self.data_dir),
+                required_exts=[".md", ".txt"],
+                recursive=False
+            ).load_data()
+            
+            logger.info(f"📄 Loaded {len(documents)} documents")
+            
+            if not documents:
+                logger.error("❌ No documents loaded. Check data directory.")
+                self.enabled = False
+                return
+            
+            # Create index
+            self.index = VectorStoreIndex.from_documents(
+                documents,
+                show_progress=True
+            )
+            
+            # Persist index
+            self.persist_dir.mkdir(parents=True, exist_ok=True)
+            self.index.storage_context.persist(persist_dir=str(self.persist_dir))
+            logger.info(f"💾 Index persisted to {self.persist_dir}")
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to create index: {e}")
+            import traceback
+            traceback.print_exc()
+            self.enabled = False
+    
+    def _create_sample_data(self):
+        """Create sample card benefit file if none exist"""
+        sample_file = self.data_dir / "sample_card.md"
+        sample_content = """# Sample Credit Card
+
+## Earning Rates
+- 4x points at restaurants
+- 4x points at U.S. supermarkets (up to $25,000/year)
+- 3x points on flights
+- 1x points on everything else
+
+## Annual Fee
+$250 (offset by $240 in credits)
+
+## Best For
+Dining and grocery spending
+
+## Important Notes
+- Supercenters like Walmart and Target do NOT count as supermarkets
+- Must activate credits to receive full value
+- No foreign transaction fees
+"""
+        sample_file.write_text(sample_content)
+        logger.info(f"📝 Created sample file: {sample_file}")
+    
+    def query_benefits(self, card_name: str, question: str) -> Optional[str]:
+        """
+        Query card benefits
+        
+        Args:
+            card_name: Name of the card
+            question: Question about the card
+        
+        Returns:
+            Answer from RAG or None
+        """
+        if not self.enabled:
+            logger.warning("RAG query skipped (not enabled)")
+            return None
+        
+        try:
+            query = f"For {card_name}: {question}"
+            logger.info(f"🔍 RAG Query: {query}")
+            response = self.query_engine.query(query)
+            return str(response)
+        except Exception as e:
+            logger.error(f"❌ Query failed: {e}")
+            import traceback
+            traceback.print_exc()
+            return None
+    
+    def get_card_context(self, card_name: str, merchant: str, category: str) -> Optional[str]:
+        """
+        Get relevant context for a card recommendation
+        
+        Args:
+            card_name: Recommended card
+            merchant: Merchant name
+            category: Spending category
+        
+        Returns:
+            Relevant context or None
+        """
+        if not self.enabled:
+            return None
+        
+        try:
+            query = f"""For {card_name} when shopping at {merchant} ({category} category):
+1. What are the earning rates for {category} purchases?
+2. Are there any spending caps or exclusions relevant to {merchant}?
+3. What are 2-3 key benefits or warnings for this type of purchase?
+
+Provide a concise summary in 2-3 sentences."""
+            
+            logger.info(f"🔍 Context Query: {card_name} at {merchant}")
+            response = self.query_engine.query(query)
+            result = str(response)
+            
+            # Clean up response
+            if len(result) > 500:
+                result = result[:500] + "..."
+            
+            return result
+            
+        except Exception as e:
+            logger.error(f"❌ Context retrieval failed: {e}")
+            return None
+    
+    def compare_cards(self, card1: str, card2: str, category: str) -> Optional[str]:
+        """
+        Compare two cards for a specific category
+        
+        Args:
+            card1: First card name
+            card2: Second card name
+            category: Spending category
+        
+        Returns:
+            Comparison or None
+        """
+        if not self.enabled:
+            return None
+        
+        try:
+            query = f"Compare {card1} vs {card2} for {category} spending. Which is better and why? Provide a concise answer in 2-3 sentences."
+            logger.info(f"🔍 Comparison: {card1} vs {card2} for {category}")
+            response = self.query_engine.query(query)
+            return str(response)
+        except Exception as e:
+            logger.error(f"❌ Comparison failed: {e}")
+            return None
+    
+    def get_spending_warnings(self, card_name: str, category: str, amount: float) -> Optional[str]:
+        """
+        Get warnings about spending caps or limitations
+        
+        Args:
+            card_name: Card name
+            category: Spending category
+            amount: Transaction amount
+        
+        Returns:
+            Warnings or None
+        """
+        if not self.enabled:
+            return None
+        
+        try:
+            query = f"For {card_name} and a ${amount:.2f} purchase in {category} category: Are there any spending caps, annual limits, or exclusions I should know about? Be specific and concise."
+            logger.info(f"🔍 Warnings: {card_name} ${amount} in {category}")
+            response = self.query_engine.query(query)
+            return str(response)
+        except Exception as e:
+            logger.error(f"❌ Warning retrieval failed: {e}")
+            return None
+
+
+# Global instance
+_rag_instance = None
+
+def get_card_benefits_rag() -> CardBenefitsRAG:
+    """Get or create the global RAG instance"""
+    global _rag_instance
+    if _rag_instance is None:
+        _rag_instance = CardBenefitsRAG()
+    return _rag_instance
+
+
+# Initialize on module import (lazy loading)
+def initialize_rag():
+    """Initialize RAG system (call this at app startup)"""
+    logger.info("🚀 Initializing LlamaIndex RAG...")
+    rag = get_card_benefits_rag()
+    if rag.enabled:
+        logger.info("✅ RAG initialized and ready")
+    else:
+        logger.warning("⚠️ RAG not available")
+    return rag