Advanced
RAG Techniques

# How to implement multi-query retrieval with LangChain
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.prompts import PromptTemplate
from typing import List
import asyncio
import os
from dotenv import load_dotenv

load_dotenv()

class AdvancedMultiQueryRetriever:
    """Learn how to build multi-query retrieval with LangChain step-by-step"""
    
    def __init__(self, vector_store, llm=None):
        self.vector_store = vector_store
        self.llm = llm or ChatGoogleGenerativeAI(
            model="gemini-2.0-flash",
            temperature=0.5,
            google_api_key=os.getenv("GOOGLE_API_KEY")
        )
        self.query_generation_prompt = PromptTemplate(
            input_variables=["question"],
            template="""You are an AI assistant helping to improve search results.
            
Generate 5 different versions of the given question to retrieve relevant documents.
Include variations that:
1. Use different keywords and synonyms
2. Are more specific or more general
3. Focus on different aspects of the question
4. Use technical vs. simple language
5. Rephrase as a statement instead of a question

Original question: {question}

Output the variations as a numbered list:
1.
2.
3.
4.
5."""
        )
    
    def generate_queries(self, original_query: str) -> List[str]:
        """Generate multiple query variations"""
        # Get LLM to generate variations
        response = self.llm.invoke(
            self.query_generation_prompt.format(question=original_query)
        )
        
        # Parse variations from response
        variations = []
        lines = response.content.strip().split('\n')
        for line in lines:
            if line.strip() and line[0].isdigit():
                # Remove number and period
                query = line.split('.', 1)[1].strip()
                if query:
                    variations.append(query)
        
        # Always include original query
        variations.insert(0, original_query)
        
        return variations[:6]  # Limit to 6 queries total
    
    async def retrieve_with_multi_query(self, query: str, k: int = 5):
        """Retrieve documents using multiple query variations"""
        # Generate query variations
        queries = self.generate_queries(query)
        print(f"Generated {len(queries)} query variations:")
        for i, q in enumerate(queries):
            print(f"  {i+1}. {q}")
        
        # Retrieve documents for each query
        all_docs = []
        unique_docs = {}
        
        for q in queries:
            docs = self.vector_store.similarity_search_with_score(q, k=k)
            for doc, score in docs:
                # Use content hash as unique identifier
                doc_id = hash(doc.page_content)
                
                if doc_id not in unique_docs:
                    unique_docs[doc_id] = {
                        "document": doc,
                        "scores": [score],
                        "matched_queries": [q]
                    }
                else:
                    unique_docs[doc_id]["scores"].append(score)
                    unique_docs[doc_id]["matched_queries"].append(q)
        
        # Aggregate scores (using max score for each doc)
        results = []
        for doc_data in unique_docs.values():
            max_score = max(doc_data["scores"])
            avg_score = sum(doc_data["scores"]) / len(doc_data["scores"])
            
            results.append({
                "document": doc_data["document"],
                "max_score": max_score,
                "avg_score": avg_score,
                "match_count": len(doc_data["matched_queries"]),
                "matched_queries": doc_data["matched_queries"]
            })
        
        # Sort by combination of max score and match count
        results.sort(
            key=lambda x: (x["match_count"], x["max_score"]), 
            reverse=True
        )
        
        return results[:k]

# Usage example - Multi-Query Retrieval Demo
print("=== Multi-Query Retrieval Demo ===")

# Create mock vector store for demonstration
class MockVectorStore:
    def similarity_search_with_score(self, query, k=5):
        # Mock search results based on query keywords
        mock_docs = [
            {"content": f"Document about {query[:20]}... with caching strategies", "score": 0.9},
            {"content": f"Advanced guide on {query[:15]}... performance optimization", "score": 0.8},
            {"content": f"Tutorial covering {query[:25]}... implementation details", "score": 0.7},
            {"content": f"Best practices for {query[:20]}... in production systems", "score": 0.6},
            {"content": f"Common patterns when {query[:30]}... troubleshooting guide", "score": 0.5},
        ]
        
        # Create mock document objects
        class MockDoc:
            def __init__(self, content):
                self.page_content = content
                self.metadata = {"source": "mock_doc.txt"}
        
        return [(MockDoc(doc["content"]), doc["score"]) for doc in mock_docs[:k]]

# Create demo multi-retriever
mock_vector_store = MockVectorStore()

# Check if API key is available
if os.getenv("GOOGLE_API_KEY"):
    # Use real LLM if API key is available
    multi_retriever = AdvancedMultiQueryRetriever(mock_vector_store)
    print("Using Google Gemini LLM for query generation")
else:
    # Fall back to mock LLM for demo purposes
    class MockLLM:
        def invoke(self, prompt):
            # Mock response for query generation
            class MockResponse:
                def __init__(self):
                    self.content = """1. What are caching strategies for web applications?
2. How to implement web application caching mechanisms?
3. Best practices for caching in web development?
4. Web app caching techniques and implementation?
5. Application-level caching for web systems?"""
            return MockResponse()
    
    mock_llm = MockLLM()
    multi_retriever = AdvancedMultiQueryRetriever(mock_vector_store, llm=mock_llm)
    print("Using mock LLM for demo (set GOOGLE_API_KEY for real LLM)")

# Demo function to run async code
async def run_multi_query_demo():
    print("\nTesting multi-query retrieval...")
    query = "How do I implement caching in a web application?"
    print(f"Original query: {query}")
    
    # Generate query variations
    variations = multi_retriever.generate_queries(query)
    print(f"\nGenerated {len(variations)} query variations:")
    for i, var in enumerate(variations, 1):
        print(f"  {i}. {var}")
    
    # Retrieve with multi-query
    results = await multi_retriever.retrieve_with_multi_query(query, k=3)
    
    print(f"\nTop {len(results)} results:")
    for i, result in enumerate(results):
        print(f"\nResult {i+1}:")
        print(f"  Score: {result['max_score']:.3f}")
        print(f"  Matched by {result['match_count']} queries")
        print(f"  Content: {result['document'].page_content[:80]}...")
    
    return results

# Run the demo
try:
    import asyncio
    results = asyncio.run(run_multi_query_demo())
    print(f"\n✓ Successfully retrieved {len(results)} documents using multi-query approach")
except Exception as e:
    print(f"Demo mode - Multi-query retrieval concept demonstrated")
    print("Note: This would work with a real vector store and LLM")

# How to implement parent-child chunking with LangChain
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.storage import InMemoryStore
from langchain.retrievers import ParentDocumentRetriever
from langchain_community.document_loaders import TextLoader
from langchain_chroma import Chroma
from langchain_google_genai import GoogleGenerativeAIEmbeddings
import uuid
import os
from dotenv import load_dotenv

load_dotenv()

class HierarchicalChunker:
    """Learn how to implement parent-child chunking in LangChain step-by-step"""
    
    def __init__(self, vector_store, docstore=None):
        self.vector_store = vector_store
        self.docstore = docstore or InMemoryStore()
        self.child_chunks = []
        self.parent_chunks = []
        
        # Child splitter - small chunks for retrieval
        self.child_splitter = RecursiveCharacterTextSplitter(
            chunk_size=400,
            chunk_overlap=50,
            separators=["\n\n", "\n", ". ", ", ", " ", ""]
        )
        
        # Parent splitter - larger chunks for context
        self.parent_splitter = RecursiveCharacterTextSplitter(
            chunk_size=2000,
            chunk_overlap=200,
            separators=["\n\n\n", "\n\n", "\n", ". "]
        )
    
    def create_hierarchical_chunks(self, documents):
        """Create parent and child chunks with relationships"""
        all_parents = []
        all_children = []
        
        for doc in documents:
            # Create parent chunks
            parent_chunks = self.parent_splitter.split_documents([doc])
            
            for parent in parent_chunks:
                # Generate unique parent ID
                parent_id = str(uuid.uuid4())
                parent.metadata["doc_id"] = parent_id
                parent.metadata["type"] = "parent"
                
                # Create child chunks from this parent
                children = self.child_splitter.split_documents([parent])
                
                for child in children:
                    # Link child to parent
                    child.metadata["parent_id"] = parent_id
                    child.metadata["type"] = "child"
                    all_children.append(child)
                
                all_parents.append(parent)
        
        # Store chunks as instance attributes
        self.parent_chunks = all_parents
        self.child_chunks = all_children
        
        return all_parents, all_children
    
    def setup_parent_child_retriever(self, documents):
        """Setup retriever with parent-child relationships"""
        # Create hierarchical chunks
        parents, children = self.create_hierarchical_chunks(documents)
        
        # Store parents in docstore
        parent_docs = {}
        for parent in parents:
            parent_docs[parent.metadata["doc_id"]] = parent
        self.docstore.mset(list(parent_docs.items()))
        
        # Add children to vector store (these will be searched)
        self.vector_store.add_documents(children)
        
        # Create parent document retriever
        retriever = ParentDocumentRetriever(
            vectorstore=self.vector_store,
            docstore=self.docstore,
            child_splitter=self.child_splitter,
            parent_splitter=self.parent_splitter
        )
        
        return retriever
    
    def retrieve_with_context(self, query: str, k: int = 4):
        """Retrieve children but return parent documents"""
        # Search for relevant children
        child_results = self.vector_store.similarity_search_with_score(
            query, 
            k=k*2,  # Get more children to find diverse parents
            filter={"type": "child"}
        )
        
        # Get unique parent documents
        parent_ids = set()
        parent_docs = []
        
        for child, score in child_results:
            parent_id = child.metadata.get("parent_id")
            
            if parent_id and parent_id not in parent_ids:
                parent_ids.add(parent_id)
                
                # Retrieve parent document
                parent = self.docstore.mget([parent_id])[0]
                
                if parent:
                    parent_docs.append({
                        "document": parent,
                        "score": score,
                        "matched_child": child.page_content[:100]
                    })
        
        # Return top k parent documents
        return parent_docs[:k]

# Create embeddings and vector store
embeddings = GoogleGenerativeAIEmbeddings(
    model="models/embedding-001",
    google_api_key=os.getenv("GOOGLE_API_KEY")
)

vector_store = Chroma(
    embedding_function=embeddings,
    persist_directory="./hierarchical_db"
)

# Example usage - Parent-Child Chunking Demo
print("=== Parent-Child Chunking Demo ===")
hierarchical_chunker = HierarchicalChunker(vector_store)

# Create sample document for testing
sample_document_content = """
# Machine Learning Fundamentals

Machine learning is a subset of artificial intelligence that enables computers to learn and make decisions from data without being explicitly programmed for every task.

## Types of Machine Learning

### Supervised Learning
Supervised learning uses labeled training data to learn a mapping function from input variables to output variables. Common algorithms include linear regression, decision trees, and neural networks.

Examples of supervised learning:
- Image classification: Training a model to recognize cats vs. dogs using labeled images
- Email spam detection: Using labeled emails to identify spam
- Stock price prediction: Using historical data to predict future prices

### Unsupervised Learning
Unsupervised learning finds hidden patterns in data without labeled examples. It includes clustering, dimensionality reduction, and association rule learning.

Examples of unsupervised learning:
- Customer segmentation: Grouping customers based on purchasing behavior
- Anomaly detection: Identifying unusual patterns in network traffic
- Market basket analysis: Finding products frequently bought together

### Reinforcement Learning
Reinforcement learning trains agents to make decisions through interaction with an environment, learning from rewards and penalties.

Examples of reinforcement learning:
- Game playing: AlphaGo mastering the game of Go
- Autonomous vehicles: Learning to navigate traffic
- Recommendation systems: Optimizing user engagement

## Key Concepts

### Training and Testing
Machine learning models are trained on a training set and evaluated on a separate test set to measure generalization performance.

### Overfitting and Underfitting
Overfitting occurs when a model learns the training data too well, including noise. Underfitting happens when a model is too simple to capture underlying patterns.

### Feature Engineering
The process of selecting, modifying, or creating input features to improve model performance.
"""

# Create sample document
class MockDocument:
    def __init__(self, content, metadata=None):
        self.page_content = content
        self.metadata = metadata or {"source": "ml_tutorial.txt"}

sample_doc = MockDocument(sample_document_content)
documents = [sample_doc]

# Setup parent-child retrieval
print("Setting up parent-child retrieval system...")
retriever = hierarchical_chunker.setup_parent_child_retriever(documents)
print(f"✓ Created {len(hierarchical_chunker.child_chunks)} child chunks")
print(f"✓ Created {len(hierarchical_chunker.parent_chunks)} parent chunks")

# Test queries
test_queries = [
    "What is supervised learning?",
    "Give me examples of reinforcement learning",
    "What is the difference between overfitting and underfitting?"
]

for query in test_queries:
    print(f"\n--- Query: '{query}' ---")
    
    # Retrieve with full parent context
    results = hierarchical_chunker.retrieve_with_context(query, k=2)
    
    print(f"Retrieved {len(results)} parent documents:")
    for i, result in enumerate(results):
        print(f"\nParent {i+1} (score: {result['score']:.3f}):")
        print(f"  Matched child: {result['matched_child'][:100]}...")
        print(f"  Parent context ({len(result['document'].page_content)} chars):")
        print(f"  {result['document'].page_content[:150]}...")

print("\n=== Chunking Strategy Comparison ===")
print(f"Child chunks: {len(hierarchical_chunker.child_chunks)} (avg ~{sum(len(c.page_content) for c in hierarchical_chunker.child_chunks)//len(hierarchical_chunker.child_chunks)} chars)")
print(f"Parent chunks: {len(hierarchical_chunker.parent_chunks)} (avg ~{sum(len(p.page_content) for p in hierarchical_chunker.parent_chunks)//len(hierarchical_chunker.parent_chunks)} chars)")
print("\n✓ Child chunks provide precise matching")
print("✓ Parent chunks provide rich context for generation")

from datetime import datetime, timedelta
from typing import Dict, Any, List
import json

class MetadataFilteredRetriever:
    """Learn how to implement metadata filtering in LangChain step-by-step"""
    
    def __init__(self, vector_store):
        self.vector_store = vector_store
        self.metadata_index = {}
        
    def build_metadata_index(self, documents):
        """Build index of metadata for fast filtering"""
        for doc in documents:
            for key, value in doc.metadata.items():
                if key not in self.metadata_index:
                    self.metadata_index[key] = set()
                self.metadata_index[key].add(value)
        
        print(f"Built metadata index with {len(self.metadata_index)} fields")
        for key, values in self.metadata_index.items():
            print(f"  {key}: {len(values)} unique values")
    
    def create_smart_filter(self, 
                          query: str, 
                          user_context: Dict[str, Any]) -> Dict:
        """Create intelligent filters based on query and context"""
        filters = {}
        
        # Time-based filtering
        if "recent" in query.lower() or "latest" in query.lower():
            week_ago = (datetime.now() - timedelta(days=7)).isoformat()
            filters["date"] = {"$gte": week_ago}
        
        # Department/team filtering based on user context
        if user_department := user_context.get("department"):
            filters["department"] = {"$in": [user_department, "company-wide"]}
        
        # Document type inference from query
        query_lower = query.lower()
        if any(word in query_lower for word in ["policy", "procedure", "guideline"]):
            filters["doc_type"] = {"$in": ["policy", "procedure"]}
        elif any(word in query_lower for word in ["report", "analysis", "study"]):
            filters["doc_type"] = {"$in": ["report", "research"]}
        
        # Access level filtering
        user_clearance = user_context.get("clearance_level", 1)
        filters["required_clearance"] = {"$lte": user_clearance}
        
        return filters
    
    def retrieve_with_metadata(self, 
                              query: str,
                              filters: Dict = None,
                              k: int = 5,
                              score_threshold: float = 0.7):
        """Retrieve with complex metadata filtering"""
        # Build search kwargs
        search_kwargs = {
            "k": k * 2,  # Get extra results for post-filtering
        }
        
        if filters:
            # Chroma expects filters to be wrapped in $and for multiple conditions
            if len(filters) > 1:
                search_kwargs["filter"] = {"$and": [
                    {key: value} for key, value in filters.items()
                ]}
            else:
                search_kwargs["filter"] = filters
        
        # Perform search
        results = self.vector_store.similarity_search_with_score(
            query,
            **search_kwargs
        )
        
        # Post-process results
        filtered_results = []
        for doc, score in results:
            if score >= score_threshold:
                # Calculate metadata relevance boost
                metadata_boost = self._calculate_metadata_boost(
                    doc.metadata, 
                    query
                )
                
                adjusted_score = score * (1 + metadata_boost)
                
                filtered_results.append({
                    "document": doc,
                    "base_score": score,
                    "adjusted_score": adjusted_score,
                    "metadata": doc.metadata
                })
        
        # Sort by adjusted score
        filtered_results.sort(
            key=lambda x: x["adjusted_score"], 
            reverse=True
        )
        
        return filtered_results[:k]
    
    def _calculate_metadata_boost(self, 
                                 metadata: Dict, 
                                 query: str) -> float:
        """Calculate relevance boost based on metadata"""
        boost = 0.0
        query_lower = query.lower()
        
        # Boost for query terms in title
        if title := metadata.get("title", "").lower():
            query_words = query_lower.split()
            matches = sum(1 for word in query_words if word in title)
            boost += matches * 0.1
        
        # Boost for recent documents
        if date_str := metadata.get("date"):
            try:
                doc_date = datetime.fromisoformat(date_str)
                days_old = (datetime.now() - doc_date).days
                if days_old < 30:
                    boost += 0.2
                elif days_old < 90:
                    boost += 0.1
            except:
                pass
        
        # Boost for high-priority documents
        if metadata.get("priority") == "high":
            boost += 0.15
        
        return min(boost, 0.5)  # Cap boost at 50%

# Create embeddings and vector store for demo
import os
from dotenv import load_dotenv
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_chroma import Chroma

load_dotenv()

# Initialize embeddings
embeddings = GoogleGenerativeAIEmbeddings(
    model="models/embedding-001",
    google_api_key=os.getenv("GOOGLE_API_KEY")
)

# Create vector store
vector_store = Chroma(
    embedding_function=embeddings,
    persist_directory="./metadata_db"
)

# Example usage - Metadata Filtering Demo
print("=== Metadata Filtering Demo ===")
metadata_retriever = MetadataFilteredRetriever(vector_store)

# User context
user_context = {
    "department": "engineering",
    "clearance_level": 3,
    "role": "senior_developer"
}

# Create smart filters
filters = metadata_retriever.create_smart_filter(
    "recent security policies for API development",
    user_context
)

print("Generated filters:", json.dumps(filters, indent=2))

# For demo purposes, we'll show the filter format instead of actual retrieval
print("\n--- Chroma Filter Format ---")
if len(filters) > 1:
    chroma_filter = {"$and": [
        {key: value} for key, value in filters.items()
    ]}
else:
    chroma_filter = filters
    
print(f"Chroma-compatible filter: {json.dumps(chroma_filter, indent=2)}")

print(f"\nUser context: {user_context}")

# Create smart filters
filters = metadata_retriever.create_smart_filter(
    "recent security policies for API development",
    user_context
)

print(f"\nGenerated filters: {json.dumps(filters, indent=2)}")

# Demo metadata-based relevance scoring
print("\n--- Metadata Boost Calculation Demo ---")
sample_docs = [
    {"content": "API Security Best Practices for 2024", "metadata": {"date": "2024-01-15", "department": "engineering", "clearance": 3}},
    {"content": "Basic API Guidelines", "metadata": {"date": "2023-06-01", "department": "general", "clearance": 1}},
    {"content": "Advanced Security Protocols", "metadata": {"date": "2024-02-20", "department": "security", "clearance": 4}},
]

for doc in sample_docs:
    # For demo, use a sample query instead of user_context
    boost = metadata_retriever._calculate_metadata_boost(doc["metadata"], "API security policies")
    
    # Also calculate user context-based boost manually
    context_boost = 0.0
    if doc["metadata"].get("department") == user_context.get("department"):
        context_boost += 0.3
    if doc["metadata"].get("clearance", 0) <= user_context.get("clearance_level", 0):
        context_boost += 0.2
    
    print(f"\nDocument: {doc['content']}")
    print(f"  Metadata: {doc['metadata']}")
    print(f"  Query-based boost: {boost:.3f}")
    print(f"  Context-based boost: {context_boost:.3f}")
    print(f"  Total boost: {(boost + context_boost):.3f}")

print("\n✓ Documents matching user's department and clearance level get higher scores")
print("✓ Recent documents receive additional boost")
print("✓ This ensures contextually relevant results for each user")

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import (
    LLMChainExtractor,
    CrossEncoderReranker
)
from langchain.prompts import PromptTemplate
from sentence_transformers import CrossEncoder
import numpy as np
from typing import List, Dict, Tuple
import os
from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI

load_dotenv()

class AdvancedReranker:
    """Multi-stage reranking with fusion capabilities"""
    
    def __init__(self, llm=None):
        self.llm = llm or ChatGoogleGenerativeAI(
            model="gemini-2.0-flash",
            temperature=0
        )
        
        # Initialize cross-encoder for reranking
        self.cross_encoder = CrossEncoder(
            'cross-encoder/ms-marco-MiniLM-L-6-v2'
        )
        
    def llm_rerank(self, query: str, documents: List, top_k: int = 5):
        """Use LLM to rerank documents based on relevance"""
        rerank_prompt = PromptTemplate(
            input_variables=["query", "document", "index"],
            template="""Given the query and document, rate the relevance on a scale of 0-10.
            
Query: {query}

Document {index}: {document}

Consider:
1. Direct answer to the query
2. Relevant background information
3. Quality and clarity of information

Relevance Score (0-10):"""
        )
        
        scored_docs = []
        
        for i, doc in enumerate(documents):
            # Get LLM relevance score
            response = self.llm.invoke(
                rerank_prompt.format(
                    query=query,
                    document=doc.page_content[:500],
                    index=i+1
                )
            )
            
            try:
                score = float(response.content.strip())
            except:
                score = 0.0
            
            scored_docs.append({
                "document": doc,
                "llm_score": score / 10.0  # Normalize to 0-1
            })
        
        # Sort by LLM score
        scored_docs.sort(key=lambda x: x["llm_score"], reverse=True)
        
        return scored_docs[:top_k]
    
    def cross_encoder_rerank(self, 
                            query: str, 
                            documents: List,
                            top_k: int = 5):
        """Use cross-encoder for precise reranking"""
        # Prepare pairs for cross-encoder
        pairs = [[query, doc.page_content] for doc in documents]
        
        # Get cross-encoder scores
        scores = self.cross_encoder.predict(pairs)
        
        # Create scored documents
        scored_docs = []
        for doc, score in zip(documents, scores):
            scored_docs.append({
                "document": doc,
                "cross_encoder_score": score
            })
        
        # Sort by score
        scored_docs.sort(
            key=lambda x: x["cross_encoder_score"], 
            reverse=True
        )
        
        return scored_docs[:top_k]
    
    def reciprocal_rank_fusion(self, 
                              result_sets: List[List[Dict]], 
                              k: int = 60):
        """Fuse multiple result sets using RRF"""
        # Track document scores
        doc_scores = {}
        
        for results in result_sets:
            for rank, result in enumerate(results):
                # Get document identifier
                doc_id = hash(result["document"].page_content)
                
                # Calculate RRF score
                rrf_score = 1.0 / (k + rank + 1)
                
                if doc_id not in doc_scores:
                    doc_scores[doc_id] = {
                        "document": result["document"],
                        "scores": [],
                        "methods": []
                    }
                
                doc_scores[doc_id]["scores"].append(rrf_score)
                doc_scores[doc_id]["methods"].append(
                    result.get("method", "unknown")
                )
        
        # Calculate final scores
        fused_results = []
        for doc_data in doc_scores.values():
            final_score = sum(doc_data["scores"])
            
            fused_results.append({
                "document": doc_data["document"],
                "fusion_score": final_score,
                "num_methods": len(doc_data["scores"]),
                "methods": list(set(doc_data["methods"]))
            })
        
        # Sort by fusion score
        fused_results.sort(
            key=lambda x: x["fusion_score"], 
            reverse=True
        )
        
        return fused_results
    
    def hybrid_rerank_and_fuse(self, 
                              query: str,
                              initial_results: List,
                              top_k: int = 5):
        """Complete reranking and fusion pipeline"""
        # Stage 1: Cross-encoder reranking
        cross_encoder_results = self.cross_encoder_rerank(
            query, 
            [r["document"] for r in initial_results],
            top_k=top_k*2
        )
        
        # Add method tag
        for r in cross_encoder_results:
            r["method"] = "cross_encoder"
        
        # Stage 2: LLM reranking (on top results only)
        top_docs = [r["document"] for r in cross_encoder_results[:10]]
        llm_results = self.llm_rerank(query, top_docs, top_k=top_k)
        
        # Add method tag
        for r in llm_results:
            r["method"] = "llm_rerank"
        
        # Stage 3: Fuse results
        fused_results = self.reciprocal_rank_fusion(
            [cross_encoder_results, llm_results],
            k=60
        )
        
        return fused_results[:top_k]

# Create embeddings and vector store for demo
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_chroma import Chroma

embeddings = GoogleGenerativeAIEmbeddings(
    model="models/embedding-001",
    google_api_key=os.getenv("GOOGLE_API_KEY")
)

vector_store = Chroma(
    embedding_function=embeddings,
    persist_directory="./rerank_db"
)

# Example usage - Reranking Demo
print("=== Advanced Reranking Demo ===")
reranker = AdvancedReranker()

# Create mock initial results for demo
class MockDoc:
    def __init__(self, content):
        self.page_content = content
        self.metadata = {"source": "database_guide.pdf"}

# Simulate initial retrieval results with varying relevance
mock_documents = [
    ("Database query optimization basics", 0.82),
    ("Advanced SQL performance tuning", 0.78),
    ("NoSQL database design patterns", 0.75),
    ("Query execution plans explained", 0.73),
    ("Database indexing strategies", 0.71),
    ("Caching strategies for databases", 0.69),
    ("Database connection pooling", 0.67),
    ("Query optimization in PostgreSQL", 0.85),  # High relevance but lower initial rank
    ("MongoDB query performance tips", 0.65),
    ("Database sharding techniques", 0.63),
]

initial_results = [(MockDoc(content), score) for content, score in mock_documents]

print(f"Initial retrieval: {len(initial_results)} documents")
print("\nTop 5 before reranking:")
for i, (doc, score) in enumerate(initial_results[:5]):
    print(f"{i+1}. [{score:.3f}] {doc.page_content}")

# Convert to expected format
formatted_results = [
    {"document": doc, "initial_score": score}
    for doc, score in initial_results
]

# Apply hybrid reranking and fusion
final_results = reranker.hybrid_rerank_and_fuse(
    query="How to optimize database queries?",
    initial_results=formatted_results,
    top_k=5
)

print("\n\nTop 5 after reranking:")
for i, result in enumerate(final_results):
    print(f"\nResult {i+1}:")
    print(f"  Fusion Score: {result['fusion_score']:.3f}")
    print(f"  Methods: {', '.join(result['methods'])}")
    print(f"  Content: {result['document'].page_content}")
    
print("\n✓ Cross-encoder reranked based on semantic relevance")
print("✓ LLM scored based on query-specific criteria")
print("✓ RRF combined both methods for optimal ranking")

from typing import Dict, List, Optional, Any
import asyncio
from datetime import datetime
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.prompts import PromptTemplate
import os
from dotenv import load_dotenv

load_dotenv()

# Note: This complete pipeline requires the following classes from previous sections:
# - AdvancedMultiQueryRetriever (from Multi-Query section above)
# - HierarchicalChunker (from Parent-Child Chunking section above)
# - MetadataFilteredRetriever (from Metadata Filtering section above)
# - AdvancedReranker (from Reranking section above)
#
# To run this demo, either:
# 1. Copy all the class definitions from the sections above into your script
# 2. Or save each class in separate files and import them
#
# For demonstration purposes, we'll create a simplified version below

class SimplifiedAdvancedRAG:
    """Learn how to build production-ready RAG with LangChain - complete implementation"""
    
    def __init__(self, vector_store, llm=None):
        self.vector_store = vector_store
        self.llm = llm or ChatGoogleGenerativeAI(
            model="gemini-2.0-flash",
            temperature=0.3,
            google_api_key=os.getenv("GOOGLE_API_KEY")
        )
        
        # Pipeline configuration
        self.pipeline_steps = []
        
    async def query(self, 
                   query: str,
                   user_context: Dict = None,
                   use_multi_query: bool = True,
                   use_reranking: bool = True,
                   k: int = 5):
        """Execute advanced RAG pipeline"""
        
        print(f"Processing query: {query}")
        
        # Step 1: Multi-query expansion
        if use_multi_query:
            print("\n1. Generating query variations...")
            initial_results = await self.multi_query.retrieve_with_multi_query(
                query, 
                k=k*3
            )
        else:
            initial_results = self.vector_store.similarity_search_with_score(
                query,
                k=k*3
            )
        
        # Step 2: Metadata filtering
        if user_context:
            print("\n2. Applying metadata filters...")
            filters = self.metadata_filter.create_smart_filter(
                query, 
                user_context
            )
            
            # Re-search with filters
            filtered_results = self.metadata_filter.retrieve_with_metadata(
                query,
                filters=filters,
                k=k*2
            )
        else:
            filtered_results = [
                {"document": r["document"], "adjusted_score": r.get("max_score", 0)}
                for r in initial_results
            ]
        
        # Step 3: Get parent documents for better context
        print("\n3. Retrieving parent contexts...")
        parent_docs = []
        for result in filtered_results[:k*2]:
            # Check if document has parent
            if parent_id := result["document"].metadata.get("parent_id"):
                parent = self.hierarchical.docstore.mget([parent_id])[0]
                if parent:
                    parent_docs.append({
                        "document": parent,
                        "child_score": result["adjusted_score"]
                    })
            else:
                parent_docs.append({
                    "document": result["document"],
                    "child_score": result["adjusted_score"]
                })
        
        # Step 4: Reranking and fusion
        if use_reranking and len(parent_docs) > k:
            print("\n4. Reranking results...")
            final_results = self.reranker.hybrid_rerank_and_fuse(
                query,
                parent_docs,
                top_k=k
            )
        else:
            final_results = parent_docs[:k]
        
        # Step 5: Generate response
        print("\n5. Generating response...")
        response = self._generate_response(query, final_results)
        
        return {
            "answer": response,
            "sources": final_results,
            "metadata": {
                "total_retrieved": len(initial_results),
                "after_filtering": len(filtered_results),
                "final_sources": len(final_results)
            }
        }
    
    def _generate_response(self, query: str, results: List[Dict]):
        """Generate response using retrieved documents"""
        # Prepare context
        context_parts = []
        for i, result in enumerate(results):
            doc = result["document"]
            source = doc.metadata.get("source", "Unknown")
            
            context_parts.append(
                f"[Source {i+1}: {source}]\n{doc.page_content}\n"
            )
        
        context = "\n---\n".join(context_parts)
        
        # Generate response
        prompt = f"""Answer the question based on the provided context. 
If the answer cannot be found in the context, say so.

Context:
{context}

Question: {query}

Answer:"""
        
        response = self.llm.invoke(prompt)
        return response.content

# Demonstration - Complete Advanced RAG Pipeline
print("\n=== Complete Advanced RAG Pipeline Demo ===")

# Initialize components (using previous sections or mock implementations)
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_chroma import Chroma
import os
from dotenv import load_dotenv

load_dotenv()

# Create embeddings and vector store
embeddings = GoogleGenerativeAIEmbeddings(
    model="models/embedding-001",
    google_api_key=os.getenv("GOOGLE_API_KEY")
)

vector_store = Chroma(
    embedding_function=embeddings,
    persist_directory="./advanced_rag_db"
)

# Create the simplified pipeline
simplified_rag = SimplifiedAdvancedRAG(vector_store)

# Demo function
async def run_advanced_rag_demo():
    # Add sample documents to demonstrate the pipeline
    sample_docs = [
        "API security best practices include using OAuth 2.0 for authentication, implementing rate limiting, and validating all input data.",
        "Modern API development requires careful attention to security headers, CORS policies, and encryption of data in transit.",
        "Zero-trust architecture is becoming the standard for API security, requiring authentication at every level."
    ]
    
    # Add documents to vector store if needed
    if hasattr(vector_store, '_collection') and vector_store._collection.count() == 0:
        print("Adding sample documents to vector store...")
        vector_store.add_texts(sample_docs)
    
    # Run the query
    result = await simplified_rag.query(
        query="What are the best practices for API security?",
        k=3
    )
    
    print("\n🎯 Query Results:")
    print(f"\nAnswer preview: {result['answer'][:200]}...")
    print(f"\nPipeline Info:")
    for key, value in result['pipeline_info'].items():
        if isinstance(value, list):
            print(f"  {key}: {len(value)} steps")
        else:
            print(f"  {key}: {value}")
    
    return result

# Run the demo
print("\nRunning complete advanced RAG pipeline...")
try:
    import asyncio
    result = asyncio.run(run_advanced_rag_demo())
    print("\n✅ Advanced RAG pipeline completed successfully!")
    print("\n👉 Note: To use all features, copy the class implementations from the sections above.")
except Exception as e:
    print(f"\n⚠️ Demo encountered an error: {type(e).__name__}")
    print("\n💡 To run this demo with full functionality:")
    print("   1. Copy all class definitions from the sections above")
    print("   2. Ensure you have a valid GOOGLE_API_KEY in your .env file")
    print("   3. Install all required dependencies")
    print("\n📚 The demo shows the conceptual flow of an advanced RAG pipeline.")