rag_system.py

#!/usr/bin/env python3
"""
Simple RAG System for Financial Documents
Practice implementation for financial.com interview
"""

import json
import os
from typing import List, Dict, Any
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.config import Settings


class FinancialRAG:
    """Retrieval-Augmented Generation system for financial documents"""
    
    def __init__(self, collection_name: str = "financial_docs"):
        """Initialize the RAG system"""
        print("🚀 Initializing Financial RAG System...")
        
        # Load embedding model (downloads ~80MB first time)
        print("📥 Loading embedding model...")
        self.embedder = SentenceTransformer('all-MiniLM-L6-v2')
        print("✅ Embedding model loaded")
        
        # Initialize ChromaDB (in-memory for practice)
        self.client = chromadb.Client(Settings(
            anonymized_telemetry=False,
            is_persistent=False  # Change to True for persistent storage
        ))
        
        # Create or get collection
        try:
            self.collection = self.client.create_collection(
                name=collection_name,
                metadata={"description": "Financial documents from financial.com"}
            )
            print(f"✅ Created new collection: {collection_name}")
        except:
            self.collection = self.client.get_collection(name=collection_name)
            print(f"✅ Using existing collection: {collection_name}")
    
    def load_documents(self, filepath: str) -> List[Dict[str, Any]]:
        """Load documents from JSON file"""
        print(f"\n📂 Loading documents from {filepath}...")
        
        with open(filepath, 'r', encoding='utf-8') as f:
            documents = json.load(f)
        
        print(f"✅ Loaded {len(documents)} documents")
        return documents
    
    def chunk_document(self, text: str, chunk_size: int = 400, 
                      overlap: int = 50) -> List[str]:
        """
        Split document into overlapping chunks
        
        Args:
            text: Document text
            chunk_size: Target chunk size in characters
            overlap: Overlap between chunks
            
        Returns:
            List of text chunks
        """
        chunks = []
        start = 0
        text_len = len(text)
        
        while start < text_len:
            end = min(start + chunk_size, text_len)
            chunk = text[start:end]
            
            # Try to break at sentence boundary
            if end < text_len:
                # Look for period followed by space
                last_period = chunk.rfind('. ')
                if last_period > chunk_size * 0.6:  # Only if chunk not too short
                    chunk = chunk[:last_period + 1]
                    end = start + last_period + 1
            
            chunks.append(chunk.strip())
            start = end - overlap if end < text_len else text_len
        
        return [c for c in chunks if c]  # Filter empty chunks
    
    def ingest(self, documents: List[Dict[str, Any]], 
               chunk_documents: bool = True):
        """
        Process and store documents in vector database
        
        Args:
            documents: List of document dicts with 'content' and 'metadata'
            chunk_documents: Whether to chunk large documents
        """
        print(f"\n⚙️  Processing documents for ingestion...")
        
        all_texts = []
        all_metadatas = []
        all_ids = []
        
        for doc in documents:
            content = doc.get('content', '')
            doc_id = doc.get('id', f"doc_{len(all_ids)}")
            metadata = doc.get('metadata', {})
            
            if chunk_documents and len(content) > 500:
                # Chunk large documents
                chunks = self.chunk_document(content)
                print(f"  📄 {doc_id}: Split into {len(chunks)} chunks")
                
                for i, chunk in enumerate(chunks):
                    all_texts.append(chunk)
                    all_metadatas.append({
                        **metadata,
                        'source_doc': doc_id,
                        'chunk_index': i,
                        'total_chunks': len(chunks)
                    })
                    all_ids.append(f"{doc_id}_chunk_{i}")
            else:
                # Keep as single document
                all_texts.append(content)
                all_metadatas.append({**metadata, 'source_doc': doc_id})
                all_ids.append(doc_id)
        
        print(f"\n🔄 Generating embeddings for {len(all_texts)} text chunks...")
        embeddings = self.embedder.encode(
            all_texts, 
            show_progress_bar=True,
            convert_to_numpy=True
        )
        
        print(f"💾 Storing in vector database...")
        self.collection.add(
            embeddings=embeddings.tolist(),
            documents=all_texts,
            metadatas=all_metadatas,
            ids=all_ids
        )
        
        print(f"✅ Successfully ingested {len(all_texts)} chunks!")
    
    def search(self, query: str, top_k: int = 3, 
               filter_metadata: Dict = None) -> List[Dict[str, Any]]:
        """
        Search for relevant documents
        
        Args:
            query: Search query
            top_k: Number of results
            filter_metadata: Optional metadata filters
            
        Returns:
            List of results with content, metadata, and scores
        """
        print(f"\n🔍 Searching for: '{query}'")
        
        # Generate query embedding
        query_embedding = self.embedder.encode([query])[0]
        
        # Search ChromaDB
        results = self.collection.query(
            query_embeddings=[query_embedding.tolist()],
            n_results=top_k,
            where=filter_metadata  # Optional filtering
        )
        
        # Format results
        formatted_results = []
        for i in range(len(results['documents'][0])):
            formatted_results.append({
                'content': results['documents'][0][i],
                'metadata': results['metadatas'][0][i],
                'distance': results['distances'][0][i],
                'similarity': 1 - results['distances'][0][i],  # Convert distance to similarity
                'id': results['ids'][0][i]
            })
        
        print(f"✅ Found {len(formatted_results)} relevant documents")
        return formatted_results
    
    def generate_answer(self, query: str, context: List[Dict[str, Any]], 
                       max_context_length: int = 1000) -> str:
        """
        Generate answer from retrieved context
        
        In a real system, this would call an LLM API.
        For practice, we'll create a structured response.
        
        Args:
            query: User question
            context: Retrieved documents
            max_context_length: Max chars to include in context
            
        Returns:
            Generated answer
        """
        # Build context from top results
        context_text = ""
        for i, doc in enumerate(context, 1):
            snippet = doc['content'][:300] + "..." if len(doc['content']) > 300 else doc['content']
            context_text += f"\n[Source {i}] {snippet}\n"
            
            if len(context_text) > max_context_length:
                break
        
        # In production: Call OpenAI/Claude API with this prompt
        # For now, return structured answer
        
        answer = f"""Based on the retrieved documents:

{context_text}

**Answer to: "{query}"**

The most relevant information indicates: {context[0]['content'][:200]}...

[Note: In production, this would be a real LLM-generated answer]

**Sources:**
"""
        for i, doc in enumerate(context, 1):
            source = doc['metadata'].get('source_doc', 'Unknown')
            answer += f"\n{i}. {source} (similarity: {doc['similarity']:.2%})"
        
        return answer
    
    def query(self, question: str, top_k: int = 3) -> Dict[str, Any]:
        """
        Complete RAG pipeline: search + generate
        
        Args:
            question: User question
            top_k: Number of documents to retrieve
            
        Returns:
            Answer with sources
        """
        # Retrieve relevant documents
        results = self.search(question, top_k=top_k)
        
        # Generate answer
        answer = self.generate_answer(question, results)
        
        return {
            'question': question,
            'answer': answer,
            'sources': results,
            'num_sources': len(results)
        }
    
    def interactive_mode(self):
        """Interactive query mode"""
        print("\n" + "="*60)
        print("📊 Financial RAG System - Interactive Mode")
        print("="*60)
        print("Type your questions (or 'quit' to exit)\n")
        
        while True:
            try:
                query = input("💬 Your question: ").strip()
                
                if query.lower() in ['quit', 'exit', 'q']:
                    print("👋 Goodbye!")
                    break
                
                if not query:
                    continue
                
                result = self.query(query, top_k=3)
                print(f"\n{result['answer']}\n")
                print("-"*60)
                
            except KeyboardInterrupt:
                print("\n👋 Goodbye!")
                break
            except Exception as e:
                print(f"❌ Error: {e}")


def demo_basic_usage():
    """Demonstrate basic RAG functionality"""
    print("\n" + "="*60)
    print("🎯 DEMO: Basic RAG System for Financial.com")
    print("="*60)
    
    # Initialize system
    rag = FinancialRAG()
    
    # Load documents
    documents = rag.load_documents('financial_data.json')
    
    # Ingest into vector DB
    rag.ingest(documents, chunk_documents=True)
    
    # Test queries
    test_queries = [
        "What is financial.com's Q4 2024 revenue?",
        "Tell me about the Kubernetes migration",
        "What AI technologies are being used?",
        "What are clients saying about the platform?",
        "How does the deployment pipeline work?",
        "What market trends were observed in January 2025?"
    ]
    
    print("\n" + "="*60)
    print("📝 TESTING QUERIES")
    print("="*60)
    
    for query in test_queries:
        print(f"\n{'='*60}")
        result = rag.query(query, top_k=2)
        print(result['answer'])
        print(f"\n✓ Retrieved {result['num_sources']} sources")
    
    # Start interactive mode
    print("\n" + "="*60)
    rag.interactive_mode()


def demo_filtered_search():
    """Demonstrate metadata filtering"""
    print("\n" + "="*60)
    print("🎯 DEMO: Filtered Search")
    print("="*60)
    
    rag = FinancialRAG(collection_name="filtered_demo")
    documents = rag.load_documents('financial_data.json')
    rag.ingest(documents, chunk_documents=False)
    
    # Search only in tech-related documents
    print("\n🔍 Searching only in 'engineering' department:")
    results = rag.search(
        "infrastructure improvements",
        top_k=3,
        filter_metadata={"department": "engineering"}
    )
    
    for i, result in enumerate(results, 1):
        print(f"\n{i}. {result['id']} (similarity: {result['similarity']:.2%})")
        print(f"   {result['content'][:150]}...")


def demo_evaluation():
    """Demonstrate simple evaluation"""
    print("\n" + "="*60)
    print("🎯 DEMO: Retrieval Evaluation")
    print("="*60)
    
    rag = FinancialRAG(collection_name="eval_demo")
    documents = rag.load_documents('financial_data.json')
    rag.ingest(documents, chunk_documents=True)
    
    # Test cases with expected results
    test_cases = [
        {
            'query': 'Q4 revenue numbers',
            'expected_source': 'earnings_q4_2024',
            'expected_in_top_k': 1
        },
        {
            'query': 'ArgoCD deployment',
            'expected_source': 'tech_infrastructure_2025',
            'expected_in_top_k': 1
        },
        {
            'query': 'machine learning initiatives',
            'expected_source': 'ml_initiatives_2025',
            'expected_in_top_k': 1
        }
    ]
    
    correct = 0
    total = len(test_cases)
    
    for test in test_cases:
        results = rag.search(test['query'], top_k=3)
        
        # Check if expected document in top results
        sources = [r['metadata'].get('source_doc', '') for r in results]
        
        if test['expected_source'] in sources[:test['expected_in_top_k']]:
            correct += 1
            status = "✅ PASS"
        else:
            status = "❌ FAIL"
        
        print(f"\n{status} Query: '{test['query']}'")
        print(f"   Expected: {test['expected_source']}")
        print(f"   Got: {sources}")
    
    accuracy = (correct / total) * 100
    print(f"\n📊 Retrieval Accuracy: {accuracy:.1f}% ({correct}/{total})")


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1:
        if sys.argv[1] == "basic":
            demo_basic_usage()
        elif sys.argv[1] == "filter":
            demo_filtered_search()
        elif sys.argv[1] == "eval":
            demo_evaluation()
        else:
            print("Usage: python rag_system.py [basic|filter|eval]")
    else:
        # Default: run basic demo
        demo_basic_usage()