haystack_demo.py

"""
Haystack RAG Demo - Working Version
Demonstrates Haystack concepts without requiring torch.
"""

import os
import re
from typing import List, Dict, Any
from collections import Counter
import math


class HaystackStyleRAG:
    """
    RAG system that mimics Haystack's architecture and concepts.
    """
    
    def __init__(self, document_dir: str = "documents"):
        """Initialize the RAG system."""
        self.document_dir = document_dir
        self.document_store = []  # Simulates Haystack's document store
        self.retriever = None
        self.pipeline = None
        self.setup_pipeline()
    
    def setup_pipeline(self):
        """Set up the pipeline (simulating Haystack's pipeline concept)."""
        print("🔧 Setting up Haystack-style pipeline...")
        
        # Simulate Haystack's document store
        self.document_store = []
        
        # Simulate Haystack's retriever
        self.retriever = BM25Retriever()
        
        # Simulate Haystack's pipeline
        self.pipeline = RAGPipeline(self.retriever)
        
        print("✅ Haystack-style pipeline setup completed!")
    
    def load_documents(self):
        """Load documents (simulating Haystack's document loading)."""
        if not os.path.exists(self.document_dir):
            print(f"❌ Document directory '{self.document_dir}' not found!")
            return
        
        print("📚 Loading documents into Haystack-style document store...")
        
        for filename in os.listdir(self.document_dir):
            if filename.endswith('.txt'):
                file_path = os.path.join(self.document_dir, filename)
                try:
                    with open(file_path, 'r', encoding='utf-8') as file:
                        content = file.read()
                        
                        # Simulate Haystack's document format
                        document = {
                            'content': content,
                            'meta': {'filename': filename, 'source': file_path},
                            'id': len(self.document_store)
                        }
                        
                        self.document_store.append(document)
                        print(f"✓ Indexed: {filename}")
                        
                except Exception as e:
                    print(f"✗ Error loading {filename}: {e}")
        
        if self.document_store:
            print(f"✅ Document store contains {len(self.document_store)} documents!")
        else:
            print("❌ No documents found to index!")
    
    def run_pipeline(self, query: str, top_k: int = 3):
        """
        Run the Haystack-style pipeline.
        
        Args:
            query: Search query
            top_k: Number of results to return
            
        Returns:
            Pipeline results
        """
        print(f"🔄 Running Haystack pipeline for query: '{query}'")
        
        # Simulate Haystack's pipeline execution
        results = self.pipeline.run(
            query=query,
            documents=self.document_store,
            top_k=top_k
        )
        
        return results


class BM25Retriever:
    """
    BM25 Retriever (simulating Haystack's BM25Retriever).
    """
    
    def __init__(self):
        self.name = "BM25Retriever"
        print(f"✅ Initialized {self.name}")
    
    def retrieve(self, query: str, documents: List[Dict], top_k: int = 3):
        """
        Retrieve relevant documents using BM25 scoring.
        
        Args:
            query: Search query
            documents: List of documents to search
            top_k: Number of results to return
            
        Returns:
            List of relevant documents with scores
        """
        if not documents:
            return []
        
        # Tokenize query
        query_words = self.tokenize(query.lower())
        
        # Calculate BM25 scores
        scored_docs = []
        for doc in documents:
            score = self.calculate_bm25_score(query_words, doc['content'])
            if score > 0:
                scored_docs.append({
                    'document': doc,
                    'score': score
                })
        
        # Sort by score and return top_k
        scored_docs.sort(key=lambda x: x['score'], reverse=True)
        return scored_docs[:top_k]
    
    def tokenize(self, text: str) -> List[str]:
        """Simple tokenization."""
        return re.findall(r'\b\w+\b', text.lower())
    
    def calculate_bm25_score(self, query_words: List[str], content: str) -> float:
        """Calculate BM25 score for a document."""
        content_lower = content.lower()
        content_words = self.tokenize(content)
        
        if not content_words:
            return 0.0
        
        # BM25 parameters
        k1 = 1.2
        b = 0.75
        
        # Calculate document length
        doc_length = len(content_words)
        avg_doc_length = doc_length  # Simplified for demo
        
        score = 0.0
        for word in query_words:
            if word in content_words:
                tf = content_words.count(word)
                # Simplified IDF calculation
                idf = math.log(1 + 1)  # Simplified for demo
                
                # BM25 formula
                numerator = tf * (k1 + 1)
                denominator = tf + k1 * (1 - b + b * (doc_length / avg_doc_length))
                score += idf * (numerator / denominator)
        
        return score


class RAGPipeline:
    """
    RAG Pipeline (simulating Haystack's Pipeline).
    """
    
    def __init__(self, retriever):
        self.retriever = retriever
        self.name = "RAGPipeline"
        print(f"✅ Initialized {self.name}")
    
    def run(self, query: str, documents: List[Dict], top_k: int = 3):
        """
        Run the RAG pipeline.
        
        Args:
            query: Search query
            documents: List of documents
            top_k: Number of results to return
            
        Returns:
            Pipeline results
        """
        print(f"🔄 Executing {self.name}...")
        
        # Step 1: Retrieve relevant documents
        retrieved_docs = self.retriever.retrieve(query, documents, top_k)
        
        # Step 2: Generate answer
        answer = self.generate_answer(query, retrieved_docs)
        
        # Step 3: Return results
        results = {
            'query': query,
            'documents': retrieved_docs,
            'answer': answer,
            'sources': [doc['document']['meta']['filename'] for doc in retrieved_docs]
        }
        
        return results
    
    def generate_answer(self, query: str, documents: List[Dict]) -> str:
        """Generate answer from retrieved documents."""
        if not documents:
            return "I couldn't find any relevant information to answer your question."
        
        # Combine context from retrieved documents
        context_parts = []
        for doc_info in documents:
            doc = doc_info['document']
            context_parts.append(f"From {doc['meta']['filename']}:\n{doc['content'][:300]}...")
        
        context = "\n\n".join(context_parts)
        answer = f"Based on the available information:\n\n{context}"
        
        return answer


def main():
    """Main function to demonstrate the Haystack-style RAG system."""
    print("🤖 Haystack-Style RAG System Demo")
    print("=" * 50)
    
    # Initialize RAG system
    rag = HaystackStyleRAG(document_dir="documents")
    
    # Load documents
    rag.load_documents()
    
    if not rag.document_store:
        print("❌ No documents loaded. Please check your documents directory.")
        return
    
    # Demo questions
    questions = [
        "What is artificial intelligence?",
        "What are neural networks?",
        "What is machine learning?",
        "How does deep learning work?"
    ]
    
    print("\n" + "="*60)
    print("HAYSTACK-STYLE RAG DEMO")
    print("="*60)
    
    for question in questions:
        print(f"\n❓ Question: {question}")
        print("-" * 60)
        
        # Run the pipeline
        results = rag.run_pipeline(question, top_k=2)
        
        if results['documents']:
            print(f"🤖 Answer: {results['answer']}")
            print(f"📚 Sources: {', '.join(results['sources'])}")
            print(f"🔍 Retrieved {len(results['documents'])} relevant documents")
        else:
            print("❌ No relevant documents found")
        
        print("\n" + "="*60)


if __name__ == "__main__":
    main()