haystack_rag.py

"""
RAG System using Haystack Pipeline
A proper implementation using Haystack's document store and retriever.
"""

import os
from haystack import Pipeline
from haystack.document_stores import InMemoryDocumentStore
from haystack.nodes import BM25Retriever, TextConverter, PreProcessor
from haystack.utils import print_answers


class HaystackRAG:
    """
    RAG system using Haystack pipeline components.
    """
    
    def __init__(self, document_dir: str = "documents"):
        """Initialize the Haystack RAG system."""
        self.document_dir = document_dir
        self.document_store = None
        self.retriever = None
        self.pipeline = None
        self.setup_pipeline()
    
    def setup_pipeline(self):
        """Set up the Haystack pipeline."""
        try:
            print("🔧 Setting up Haystack pipeline...")
            
            # Create document store
            self.document_store = InMemoryDocumentStore(use_bm25=True)
            
            # Create retriever
            self.retriever = BM25Retriever(document_store=self.document_store)
            
            # Create pipeline
            self.pipeline = Pipeline()
            self.pipeline.add_node(component=self.retriever, name="Retriever", inputs=["Query"])
            
            print("✅ Haystack pipeline setup completed!")
            
        except Exception as e:
            print(f"❌ Error setting up pipeline: {e}")
            print("💡 Make sure Haystack is properly installed")
    
    def load_documents(self):
        """Load documents from directory."""
        if not os.path.exists(self.document_dir):
            print(f"❌ Document directory '{self.document_dir}' not found!")
            return
        
        print("📚 Loading documents...")
        documents = []
        
        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()
                        documents.append({
                            'content': content,
                            'meta': {'filename': filename, 'source': file_path}
                        })
                        print(f"✓ Loaded: {filename}")
                except Exception as e:
                    print(f"✗ Error loading {filename}: {e}")
        
        if documents:
            # Add documents to document store
            self.document_store.write_documents(documents)
            print(f"✅ Successfully indexed {len(documents)} documents!")
        else:
            print("❌ No documents found to index!")
    
    def search(self, query: str, top_k: int = 3):
        """
        Search for relevant documents using Haystack.
        
        Args:
            query: Search query
            top_k: Number of results to return
            
        Returns:
            Search results
        """
        try:
            results = self.pipeline.run(query=query, params={"Retriever": {"top_k": top_k}})
            return results
        except Exception as e:
            print(f"❌ Error during search: {e}")
            return None
    
    def ask_question(self, question: str, top_k: int = 3):
        """
        Ask a question and get an answer using Haystack RAG.
        
        Args:
            question: The question to ask
            top_k: Number of documents to retrieve
            
        Returns:
            Dictionary containing the answer and relevant documents
        """
        print(f"\n❓ Question: {question}")
        print("-" * 60)
        
        # Search for relevant documents
        results = self.search(question, top_k)
        
        if not results or not results.get('documents'):
            print("❌ No relevant documents found.")
            return {
                'answer': "I couldn't find any relevant information to answer your question.",
                'documents': [],
                'sources': []
            }
        
        # Extract documents and sources
        documents = results.get('documents', [])
        sources = [doc.meta.get('filename', 'Unknown') for doc in documents]
        
        # Generate simple answer
        context_parts = []
        for doc in documents:
            context_parts.append(f"From {doc.meta.get('filename', 'Unknown')}:\n{doc.content[:300]}...")
        
        context = "\n\n".join(context_parts)
        answer = f"Based on the available information:\n\n{context}"
        
        print(f"🤖 Answer: {answer}")
        print(f"📚 Sources: {', '.join(sources)}")
        
        return {
            'answer': answer,
            'documents': documents,
            'sources': sources
        }


def main():
    """Main function to demonstrate the Haystack RAG system."""
    print("🤖 Haystack RAG System Demo")
    print("=" * 50)
    
    # Initialize RAG system
    rag = HaystackRAG(document_dir="documents")
    
    if not rag.pipeline:
        print("❌ Pipeline setup failed. Please check your Haystack installation.")
        return
    
    # Load documents
    rag.load_documents()
    
    if not rag.document_store.get_document_count():
        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?",
        "What are the applications of AI?"
    ]
    
    print("\n" + "="*60)
    print("HAYSTACK RAG DEMO")
    print("="*60)
    
    for question in questions:
        result = rag.ask_question(question)
        print("\n" + "="*60)


if __name__ == "__main__":
    main()