Building Production-Ready RAG Systems: Lessons from the Field

import fitz  # PyMuPDF
from PIL import Image
import pytesseract
from typing import Dict, List, Tuple
import logging

class DocumentExtractor:
    """Production-grade document extraction with fallback strategies."""

    def __init__(self, ocr_confidence_threshold: float = 0.7):
        self.ocr_threshold = ocr_confidence_threshold
        self.logger = logging.getLogger(__name__)

    def extract_pdf(self, pdf_path: str) -> Dict[str, any]:
        """
        Extract text from PDF with automatic strategy selection.

        Returns a dict with:
        - text: extracted content
        - method: extraction method used (text/ocr/hybrid)
        - confidence: quality score (0-1)
        - tables: list of extracted tables
        - metadata: PDF metadata
        """
        doc = fitz.open(pdf_path)

        # Try native text extraction first
        text_content = []
        low_text_pages = []

        for page_num, page in enumerate(doc):
            text = page.get_text()

            # If page has very little text, it's likely scanned
            if len(text.strip()) < 50:
                low_text_pages.append(page_num)
            else:
                text_content.append(text)

        # Use OCR for pages with little/no text
        if low_text_pages:
            self.logger.info(f"Using OCR for {len(low_text_pages)} pages")
            ocr_results = self._ocr_pages(doc, low_text_pages)
            text_content.extend(ocr_results['text'])
            confidence = ocr_results['avg_confidence']
            method = "hybrid" if text_content else "ocr"
        else:
            confidence = 1.0
            method = "text"

        # Extract tables separately
        tables = self._extract_tables(pdf_path)

        return {
            'text': '\n\n'.join(text_content),
            'method': method,
            'confidence': confidence,
            'tables': tables,
            'metadata': doc.metadata,
            'page_count': len(doc)
        }

    def _ocr_pages(self, doc, page_numbers: List[int]) -> Dict:
        """OCR specific pages and return text with confidence scores."""
        texts = []
        confidences = []

        for page_num in page_numbers:
            page = doc[page_num]
            # Render page to image at 300 DPI for good OCR quality
            pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)

            # Get OCR data with confidence scores
            ocr_data = pytesseract.image_to_data(
                img,
                output_type=pytesseract.Output.DICT
            )

            # Calculate average confidence for this page
            conf_scores = [int(c) for c in ocr_data['conf'] if c != '-1']
            avg_conf = sum(conf_scores) / len(conf_scores) if conf_scores else 0

            # Only include text if confidence is acceptable
            if avg_conf >= self.ocr_threshold * 100:
                page_text = ' '.join(ocr_data['text'])
                texts.append(page_text)
                confidences.append(avg_conf / 100)
            else:
                self.logger.warning(
                    f"Page {page_num} OCR confidence too low: {avg_conf:.2f}%"
                )

        return {
            'text': texts,
            'avg_confidence': sum(confidences) / len(confidences) if confidences else 0
        }

    def _extract_tables(self, pdf_path: str) -> List[Dict]:
        """Extract tables with structure preserved."""
        try:
            import camelot
            tables = camelot.read_pdf(pdf_path, pages='all', flavor='lattice')

            return [{
                'page': table.page,
                'data': table.df.to_dict(),
                'shape': table.shape,
                'accuracy': table.accuracy
            } for table in tables]
        except Exception as e:
            self.logger.error(f"Table extraction failed: {e}")
            return []

# Usage example
extractor = DocumentExtractor(ocr_confidence_threshold=0.7)
result = extractor.extract_pdf("complex_document.pdf")

if result['confidence'] < 0.8:
    print(f"⚠️ Low confidence extraction ({result['confidence']:.2%})")
    print(f"Method used: {result['method']}")
    # Flag for manual review in production

from bs4 import BeautifulSoup
from typing import List, Dict
import re

class HTMLExtractor:
    """Extract structured content from HTML preserving document structure."""

    def __init__(self):
        self.heading_tags = ['h1', 'h2', 'h3', 'h4', 'h5', 'h6']

    def extract_with_structure(self, html_content: str) -> Dict:
        """
        Extract HTML with preserved structure for better chunking.

        Returns sections with heading hierarchy, allowing semantic chunking
        that respects document structure rather than arbitrary splits.
        """
        soup = BeautifulSoup(html_content, 'html.parser')

        # Remove script, style, and navigation elements
        for tag in soup(['script', 'style', 'nav', 'footer', 'aside']):
            tag.decompose()

        sections = []
        current_section = {
            'heading': '',
            'level': 0,
            'content': [],
            'position': 0
        }

        for elem in soup.find_all(['p', 'li', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'table']):
            if elem.name in self.heading_tags:
                # Save previous section if it has content
                if current_section['content']:
                    sections.append(current_section)

                # Start new section
                level = int(elem.name[1])  # h1 -> 1, h2 -> 2, etc.
                current_section = {
                    'heading': elem.get_text(strip=True),
                    'level': level,
                    'content': [],
                    'position': len(sections)
                }

            elif elem.name == 'table':
                # Store tables separately with row/column structure
                current_section['content'].append({
                    'type': 'table',
                    'data': self._extract_table(elem)
                })

            elif elem.name in ['p', 'li']:
                text = elem.get_text(strip=True)
                if text:
                    # Preserve link context
                    links = elem.find_all('a')
                    link_context = [
                        {'text': a.get_text(), 'href': a.get('href')}
                        for a in links
                    ]

                    current_section['content'].append({
                        'type': 'text',
                        'text': text,
                        'links': link_context
                    })

        # Add final section
        if current_section['content']:
            sections.append(current_section)

        return {
            'sections': sections,
            'title': soup.find('title').get_text() if soup.find('title') else '',
            'metadata': self._extract_metadata(soup)
        }

    def _extract_table(self, table_elem) -> Dict:
        """Extract table with structure preserved."""
        rows = []
        for tr in table_elem.find_all('tr'):
            cells = [td.get_text(strip=True) for td in tr.find_all(['td', 'th'])]
            rows.append(cells)

        return {
            'headers': rows[0] if rows else [],
            'rows': rows[1:] if len(rows) > 1 else []
        }

    def _extract_metadata(self, soup) -> Dict:
        """Extract page metadata from meta tags."""
        metadata = {}

        # Common meta tags
        for meta in soup.find_all('meta'):
            name = meta.get('name') or meta.get('property')
            content = meta.get('content')
            if name and content:
                metadata[name] = content

        return metadata

# Usage
extractor = HTMLExtractor()
result = extractor.extract_with_structure(html_content)

# Now you can chunk by section, preserving semantic boundaries
for section in result['sections']:
    print(f"{'#' * section['level']} {section['heading']}")
    print(f"Position: {section['position']}")
    # This section stays together during chunking

from typing import List, Dict
import tiktoken
from nltk.tokenize import sent_tokenize, ParagraphTokenizer
import nltk

# Download required NLTK data
nltk.download('punkt', quiet=True)

class SemanticChunker:
    """Chunk text at semantic boundaries, not arbitrary token counts."""

    def __init__(
        self,
        model_name: str = "gpt-4",
        target_chunk_size: int = 300,
        max_chunk_size: int = 500,
        overlap_sentences: int = 2
    ):
        self.tokenizer = tiktoken.encoding_for_model(model_name)
        self.target_chunk_size = target_chunk_size
        self.max_chunk_size = max_chunk_size
        self.overlap_sentences = overlap_sentences

    def chunk_document(self, sections: List[Dict]) -> List[Dict]:
        """
        Chunk document by semantic boundaries.

        Input: Structured sections from HTMLExtractor or similar
        Output: Chunks with metadata and overlap

        Each chunk maintains:
        - The heading hierarchy (crucial for context)
        - Position in document
        - Overlapping sentences with adjacent chunks
        """
        chunks = []

        for section in sections:
            section_heading = section['heading']
            section_level = section['level']

            # Extract all text content from section
            section_text = self._extract_section_text(section)

            # Split into paragraphs first (respects document structure)
            paragraphs = section_text.split('\n\n')

            current_chunk_paras = []
            current_tokens = 0

            for para in paragraphs:
                para = para.strip()
                if not para:
                    continue

                para_tokens = len(self.tokenizer.encode(para))

                # If single paragraph exceeds max, split by sentences
                if para_tokens > self.max_chunk_size:
                    # Finish current chunk if any
                    if current_chunk_paras:
                        chunk_text = '\n\n'.join(current_chunk_paras)
                        chunks.append(self._create_chunk(
                            chunk_text, section_heading, section_level, section['position']
                        ))
                        current_chunk_paras = []
                        current_tokens = 0

                    # Split long paragraph by sentences
                    chunks.extend(
                        self._chunk_long_paragraph(
                            para, section_heading, section_level, section['position']
                        )
                    )
                    continue

                # If adding this para exceeds target, finish current chunk
                if current_tokens + para_tokens > self.target_chunk_size and current_chunk_paras:
                    chunk_text = '\n\n'.join(current_chunk_paras)
                    chunks.append(self._create_chunk(
                        chunk_text, section_heading, section_level, section['position']
                    ))
                    current_chunk_paras = []
                    current_tokens = 0

                # Add paragraph to current chunk
                current_chunk_paras.append(para)
                current_tokens += para_tokens

            # Add final chunk if any content remains
            if current_chunk_paras:
                chunk_text = '\n\n'.join(current_chunk_paras)
                chunks.append(self._create_chunk(
                    chunk_text, section_heading, section_level, section['position']
                ))

        # Add overlaps between adjacent chunks
        chunks = self._add_overlaps(chunks)

        return chunks

    def _chunk_long_paragraph(
        self, paragraph: str, heading: str, level: int, position: int
    ) -> List[Dict]:
        """Split long paragraph by sentences."""
        sentences = sent_tokenize(paragraph)
        chunks = []

        current_sentences = []
        current_tokens = 0

        for sentence in sentences:
            sent_tokens = len(self.tokenizer.encode(sentence))

            if current_tokens + sent_tokens > self.target_chunk_size and current_sentences:
                chunk_text = ' '.join(current_sentences)
                chunks.append(self._create_chunk(chunk_text, heading, level, position))
                current_sentences = []
                current_tokens = 0

            current_sentences.append(sentence)
            current_tokens += sent_tokens

        if current_sentences:
            chunk_text = ' '.join(current_sentences)
            chunks.append(self._create_chunk(chunk_text, heading, level, position))

        return chunks

    def _create_chunk(self, text: str, heading: str, level: int, position: int) -> Dict:
        """Create chunk with metadata."""
        tokens = len(self.tokenizer.encode(text))

        return {
            'text': text,
            'heading': heading,
            'heading_level': level,
            'section_position': position,
            'token_count': tokens,
            'sentences': sent_tokenize(text)
        }

    def _add_overlaps(self, chunks: List[Dict]) -> List[Dict]:
        """Add sentence-level overlap between adjacent chunks."""
        for i in range(len(chunks) - 1):
            current_chunk = chunks[i]
            next_chunk = chunks[i + 1]

            # Take last N sentences from current chunk
            overlap_sentences = current_chunk['sentences'][-self.overlap_sentences:]

            # Add to next chunk's beginning
            next_chunk['overlap_prev'] = ' '.join(overlap_sentences)
            next_chunk['text'] = next_chunk['overlap_prev'] + '\n\n' + next_chunk['text']

        return chunks

    def _extract_section_text(self, section: Dict) -> str:
        """Extract plain text from structured section."""
        texts = []
        for item in section.get('content', []):
            if item['type'] == 'text':
                texts.append(item['text'])
            elif item['type'] == 'table':
                # Convert table to text representation
                table_text = self._table_to_text(item['data'])
                texts.append(table_text)
        return '\n\n'.join(texts)

    def _table_to_text(self, table_data: Dict) -> str:
        """Convert structured table to readable text."""
        lines = []

        # Headers
        if table_data.get('headers'):
            lines.append(' | '.join(table_data['headers']))
            lines.append('-' * 40)

        # Rows
        for row in table_data.get('rows', []):
            lines.append(' | '.join(str(cell) for cell in row))

        return '\n'.join(lines)

# Usage
chunker = SemanticChunker(
    model_name="gpt-4",
    target_chunk_size=300,
    max_chunk_size=500,
    overlap_sentences=2
)

# From our HTML extractor
html_extractor = HTMLExtractor()
structured_doc = html_extractor.extract_with_structure(html_content)

# Chunk by semantic boundaries
chunks = chunker.chunk_document(structured_doc['sections'])

for chunk in chunks:
    print(f"Heading: {chunk['heading']} (Level {chunk['heading_level']})")
    print(f"Tokens: {chunk['token_count']}")
    print(f"Text: {chunk['text'][:100]}...")
    print()

from langchain.retrievers import BM25Retriever, EnsembleRetriever
from langchain.vectorstores import Pinecone
from langchain.embeddings import OpenAIEmbeddings
from sentence_transformers import CrossEncoder
from typing import List, Dict, Tuple
import numpy as np

class HybridRetriever:
    """
    Production-grade hybrid retriever combining BM25 and vector search.

    Key design decisions:
    - BM25 for exact keyword matches (handles acronyms, proper nouns)
    - Vector search for semantic similarity (handles paraphrases, synonyms)
    - RRF for score-free fusion (no score normalization needed)
    - Cross-encoder for precise re-ranking (sees query + doc together)
    """

    def __init__(
        self,
        vector_store,
        documents: List[str],
        embeddings,
        cross_encoder_model: str = "cross-encoder/ms-marco-MiniLM-L-12-v2",
        bm25_weight: float = 0.4,
        vector_weight: float = 0.6,
        top_k_initial: int = 50,
        top_k_rerank: int = 20,
        top_k_final: int = 10
    ):
        """
        Args:
            vector_store: Vector store (Pinecone, Weaviate, etc.)
            documents: Raw documents for BM25 indexing
            embeddings: Embedding model for vector search
            cross_encoder_model: Model for re-ranking
            bm25_weight: Weight for BM25 scores (default 0.4)
            vector_weight: Weight for vector scores (default 0.6)
            top_k_initial: Candidates to retrieve from each method
            top_k_rerank: Candidates to re-rank with cross-encoder
            top_k_final: Final results to return
        """
        # Initialize BM25 retriever
        self.bm25_retriever = BM25Retriever.from_texts(
            documents,
            k=top_k_initial
        )

        # Vector store retriever
        self.vector_retriever = vector_store.as_retriever(
            search_kwargs={"k": top_k_initial}
        )

        # Cross-encoder for re-ranking
        self.cross_encoder = CrossEncoder(cross_encoder_model)

        # Parameters
        self.bm25_weight = bm25_weight
        self.vector_weight = vector_weight
        self.top_k_rerank = top_k_rerank
        self.top_k_final = top_k_final

        self.documents = documents

    def retrieve(self, query: str) -> List[Dict]:
        """
        Retrieve documents using hybrid search with re-ranking.

        Steps:
        1. BM25 retrieval (keyword-based)
        2. Vector retrieval (semantic)
        3. Reciprocal Rank Fusion
        4. Cross-encoder re-ranking
        5. Return top K

        Returns list of dicts with:
        - text: document text
        - score: final relevance score
        - rank: final ranking position
        - retrieval_method: how it was found (bm25/vector/both)
        """
        # Step 1: BM25 retrieval
        bm25_results = self.bm25_retriever.get_relevant_documents(query)

        # Step 2: Vector retrieval
        vector_results = self.vector_retriever.get_relevant_documents(query)

        # Step 3: Reciprocal Rank Fusion
        fused_results = self._reciprocal_rank_fusion(
            query, bm25_results, vector_results
        )

        # Step 4: Cross-encoder re-ranking on top candidates
        reranked_results = self._rerank_with_cross_encoder(
            query,
            fused_results[:self.top_k_rerank]
        )

        # Step 5: Return top K
        return reranked_results[:self.top_k_final]

    def _reciprocal_rank_fusion(
        self,
        query: str,
        bm25_results: List,
        vector_results: List,
        k: int = 60  # RRF constant, typically 60
    ) -> List[Dict]:
        """
        Combine rankings using Reciprocal Rank Fusion.

        RRF formula for each document d:
        RRF(d) = Σ(1 / (k + rank_i(d)))

        Where rank_i(d) is the rank of document d in retrieval method i.

        Why RRF works:
        - No score normalization needed (works with ranks only)
        - Documents appearing high in multiple rankings get boosted
        - Robust to score scale differences between methods
        """
        # Build rank dictionaries
        bm25_ranks = {doc.page_content: rank for rank, doc in enumerate(bm25_results)}
        vector_ranks = {doc.page_content: rank for rank, doc in enumerate(vector_results)}

        # All unique documents
        all_docs = set(bm25_ranks.keys()) | set(vector_ranks.keys())

        # Calculate RRF scores
        rrf_scores = {}
        for doc in all_docs:
            score = 0
            retrieval_methods = []

            if doc in bm25_ranks:
                score += self.bm25_weight / (k + bm25_ranks[doc])
                retrieval_methods.append('bm25')

            if doc in vector_ranks:
                score += self.vector_weight / (k + vector_ranks[doc])
                retrieval_methods.append('vector')

            rrf_scores[doc] = {
                'score': score,
                'retrieval_method': '+'.join(retrieval_methods),
                'bm25_rank': bm25_ranks.get(doc, None),
                'vector_rank': vector_ranks.get(doc, None)
            }

        # Sort by RRF score
        sorted_docs = sorted(
            rrf_scores.items(),
            key=lambda x: x[1]['score'],
            reverse=True
        )

        return [
            {
                'text': doc,
                'rrf_score': data['score'],
                'retrieval_method': data['retrieval_method'],
                'bm25_rank': data['bm25_rank'],
                'vector_rank': data['vector_rank']
            }
            for doc, data in sorted_docs
        ]

    def _rerank_with_cross_encoder(
        self,
        query: str,
        candidates: List[Dict]
    ) -> List[Dict]:
        """
        Re-rank candidates using cross-encoder.

        Cross-encoders see query and document together, enabling
        much more accurate relevance scoring than bi-encoders.

        Trade-off: Much slower (must run inference for each pair),
        so only use on top candidates from initial retrieval.
        """
        if not candidates:
            return []

        # Prepare query-document pairs
        pairs = [(query, cand['text']) for cand in candidates]

        # Get cross-encoder scores
        # Shape: (num_candidates,) with scores 0-1
        ce_scores = self.cross_encoder.predict(pairs)

        # Add cross-encoder scores to candidates
        for cand, score in zip(candidates, ce_scores):
            cand['cross_encoder_score'] = float(score)
            # Weighted combination of RRF and cross-encoder
            cand['final_score'] = 0.3 * cand['rrf_score'] + 0.7 * score

        # Sort by final score
        candidates.sort(key=lambda x: x['final_score'], reverse=True)

        # Add final ranks
        for rank, cand in enumerate(candidates):
            cand['rank'] = rank + 1

        return candidates


# Usage example
from langchain.vectorstores import Pinecone
from langchain.embeddings import OpenAIEmbeddings
import pinecone

# Initialize vector store
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

pinecone.init(api_key="your-key", environment="us-west1-gcp")
vector_store = Pinecone.from_documents(
    documents,
    embeddings,
    index_name="production-rag"
)

# Initialize hybrid retriever
retriever = HybridRetriever(
    vector_store=vector_store,
    documents=[doc.page_content for doc in documents],
    embeddings=embeddings,
    bm25_weight=0.4,
    vector_weight=0.6,
    top_k_final=10
)

# Retrieve
query = "What are the key findings about climate change?"
results = retriever.retrieve(query)

for result in results:
    print(f"Rank {result['rank']}: Score {result['final_score']:.3f}")
    print(f"  Method: {result['retrieval_method']}")
    print(f"  BM25 rank: {result['bm25_rank']}, Vector rank: {result['vector_rank']}")
    print(f"  Text: {result['text'][:150]}...")
    print()

from openai import OpenAI
from typing import List, Dict
import json

class RAGGenerator:
    """Generate responses from retrieved context with proper grounding."""

    def __init__(
        self,
        model: str = "gpt-4-turbo-preview",
        temperature: float = 0.0  # Low temp for factual responses
    ):
        self.client = OpenAI()
        self.model = model
        self.temperature = temperature

    def generate(
        self,
        query: str,
        retrieved_chunks: List[Dict],
        response_style: str = "concise"  # concise, detailed, technical
    ) -> Dict:
        """
        Generate answer from retrieved context.

        Args:
            query: User question
            retrieved_chunks: List of dicts with 'text', 'heading', 'score'
            response_style: Answer style (concise/detailed/technical)

        Returns:
            Dict with:
            - answer: Generated response
            - sources: List of source chunks used
            - confidence: Estimated confidence (based on retrieval scores)
        """
        # Assemble context with source IDs
        context = self._assemble_context(retrieved_chunks)

        # Select appropriate system message
        system_message = self._get_system_message(response_style)

        # Build user message
        user_message = self._build_user_message(query, context)

        # Generate response
        response = self.client.chat.completions.create(
            model=self.model,
            temperature=self.temperature,
            messages=[
                {"role": "system", "content": system_message},
                {"role": "user", "content": user_message}
            ]
        )

        answer = response.choices[0].message.content

        # Extract citations from answer
        sources = self._extract_citations(answer, retrieved_chunks)

        # Estimate confidence based on retrieval scores
        confidence = self._estimate_confidence(retrieved_chunks)

        return {
            'answer': answer,
            'sources': sources,
            'confidence': confidence,
            'model': self.model,
            'tokens_used': response.usage.total_tokens
        }

    def _get_system_message(self, style: str) -> str:
        """Get system message based on response style."""
        base_instruction = """You are a helpful assistant that answers questions using ONLY the provided context.

CRITICAL RULES:
1. Base your answer EXCLUSIVELY on the provided context
2. If the context doesn't contain enough information, say "I don't have enough information to answer this question"
3. Cite sources inline using [Source X] notation
4. Do not make up information or use external knowledge
5. If multiple sources contradict each other, mention the contradiction"""

        style_instructions = {
            "concise": "\n6. Provide concise, direct answers without unnecessary elaboration",
            "detailed": "\n6. Provide comprehensive, detailed explanations with examples",
            "technical": "\n6. Use technical language and include relevant technical details"
        }

        return base_instruction + style_instructions.get(style, "")

    def _assemble_context(self, chunks: List[Dict]) -> str:
        """
        Assemble retrieved chunks into formatted context.

        Key decisions:
        - Order by relevance (highest first) - LLMs have primacy bias
        - Include metadata (heading, source) for context
        - Add source IDs for citation
        - Use clear delimiters between chunks
        """
        context_parts = []

        for i, chunk in enumerate(chunks, 1):
            # Format: [Source X] (from "Section Heading")
            # Content here...
            source_id = f"[Source {i}]"
            heading = chunk.get('heading', 'Unknown Section')
            text = chunk['text']
            score = chunk.get('final_score', chunk.get('score', 0))

            context_parts.append(
                f"{source_id} (from \"{heading}\", relevance: {score:.2f})\n{text}"
            )

        return "\n\n---\n\n".join(context_parts)

    def _build_user_message(self, query: str, context: str) -> str:
        """Build user message with context and question."""
        return f"""Context:
{context}

---

Question: {query}

Instructions:
- Answer based ONLY on the context above
- Cite sources using [Source X] notation
- If information is insufficient, explicitly state what's missing
- If sources contradict, mention it

Answer:"""

    def _extract_citations(
        self,
        answer: str,
        chunks: List[Dict]
    ) -> List[Dict]:
        """
        Extract which sources were cited in the answer.

        This is important for:
        - Showing users where information came from
        - Validating that model actually used the context
        - Computing faithfulness metrics
        """
        import re

        # Find all [Source X] citations
        citations = re.findall(r'\[Source (\d+)\]', answer)
        cited_indices = set(int(c) - 1 for c in citations)  # Convert to 0-indexed

        sources = []
        for idx in cited_indices:
            if idx < len(chunks):
                chunk = chunks[idx]
                sources.append({
                    'source_id': idx + 1,
                    'text': chunk['text'],
                    'heading': chunk.get('heading', 'Unknown'),
                    'score': chunk.get('final_score', chunk.get('score', 0))
                })

        return sources

    def _estimate_confidence(self, chunks: List[Dict]) -> float:
        """
        Estimate answer confidence based on retrieval quality.

        Heuristics:
        - Higher retrieval scores → higher confidence
        - Multiple high-scoring chunks → higher confidence
        - Low diversity in scores → lower confidence (might be guessing)
        """
        if not chunks:
            return 0.0

        scores = [
            chunk.get('final_score', chunk.get('score', 0))
            for chunk in chunks
        ]

        # Average of top 3 scores (or all if fewer than 3)
        top_scores = sorted(scores, reverse=True)[:3]
        avg_score = sum(top_scores) / len(top_scores)

        # Confidence is roughly proportional to retrieval quality
        # Scale: 0.0-1.0
        confidence = min(avg_score, 1.0)

        return confidence


# Complete RAG Pipeline Example
class ProductionRAGSystem:
    """End-to-end RAG system combining all components."""

    def __init__(self):
        # Initialize components (assuming they're already set up)
        self.retriever = HybridRetriever(...)  # From previous code
        self.generator = RAGGenerator(model="gpt-4-turbo-preview")

    def query(
        self,
        user_query: str,
        response_style: str = "concise",
        min_confidence: float = 0.5
    ) -> Dict:
        """
        Process user query through complete RAG pipeline.

        Args:
            user_query: User's question
            response_style: Answer style (concise/detailed/technical)
            min_confidence: Minimum confidence threshold to return answer

        Returns:
            Dict with answer, sources, confidence, and metadata
        """
        # Step 1: Retrieve relevant chunks
        retrieved_chunks = self.retriever.retrieve(user_query)

        # Step 2: Generate answer
        result = self.generator.generate(
            query=user_query,
            retrieved_chunks=retrieved_chunks,
            response_style=response_style
        )

        # Step 3: Check confidence threshold
        if result['confidence'] < min_confidence:
            result['warning'] = (
                f"Low confidence ({result['confidence']:.2f}). "
                "Answer may be unreliable."
            )

        # Step 4: Log for monitoring
        self._log_query(user_query, result, retrieved_chunks)

        return result

    def _log_query(self, query: str, result: Dict, chunks: List[Dict]):
        """Log query for monitoring and improvement."""
        # In production, log to your observability system
        log_entry = {
            'query': query,
            'confidence': result['confidence'],
            'num_sources': len(result['sources']),
            'tokens_used': result['tokens_used'],
            'top_retrieval_score': chunks[0]['final_score'] if chunks else 0
        }
        # logger.info(json.dumps(log_entry))


# Usage
rag_system = ProductionRAGSystem()

result = rag_system.query(
    user_query="What are the main benefits of hybrid search in RAG systems?",
    response_style="detailed"
)

print("Answer:", result['answer'])
print("\nSources used:")
for source in result['sources']:
    print(f"  - Source {source['source_id']}: {source['heading']}")
print(f"\nConfidence: {result['confidence']:.2%}")

if 'warning' in result:
    print(f"⚠️  {result['warning']}")

from typing import List, Dict, Set
import numpy as np
from openai import OpenAI
import json

class RAGEvaluator:
    """Comprehensive evaluation framework for RAG systems."""

    def __init__(self):
        self.client = OpenAI()

    def evaluate_retrieval(
        self,
        query: str,
        retrieved_docs: List[Dict],
        relevant_doc_ids: Set[str],
        k: int = 10
    ) -> Dict:
        """
        Evaluate retrieval quality against golden dataset.

        Args:
            query: The search query
            retrieved_docs: List of retrieved documents with 'id' field
            relevant_doc_ids: Set of known relevant document IDs
            k: Number of top results to evaluate

        Returns:
            Dict with precision, recall, MRR, NDCG metrics
        """
        retrieved_ids = [doc['id'] for doc in retrieved_docs[:k]]

        # Precision@K: What fraction of retrieved docs are relevant?
        relevant_retrieved = set(retrieved_ids) & relevant_doc_ids
        precision_at_k = len(relevant_retrieved) / k if k > 0 else 0

        # Recall@K: What fraction of relevant docs were retrieved?
        recall_at_k = (
            len(relevant_retrieved) / len(relevant_doc_ids)
            if relevant_doc_ids else 0
        )

        # Mean Reciprocal Rank: Position of first relevant doc
        mrr = 0
        for i, doc_id in enumerate(retrieved_ids, 1):
            if doc_id in relevant_doc_ids:
                mrr = 1 / i
                break

        # NDCG: Ranking quality with graded relevance
        # (simplified version with binary relevance)
        dcg = sum(
            (1 if doc_id in relevant_doc_ids else 0) / np.log2(i + 1)
            for i, doc_id in enumerate(retrieved_ids, 1)
        )

        # Ideal DCG (all relevant docs at top)
        ideal_ranking = [1] * min(len(relevant_doc_ids), k)
        idcg = sum(1 / np.log2(i + 1) for i in range(1, len(ideal_ranking) + 1))

        ndcg = dcg / idcg if idcg > 0 else 0

        return {
            'precision@k': precision_at_k,
            'recall@k': recall_at_k,
            'mrr': mrr,
            'ndcg': ndcg,
            'num_relevant_retrieved': len(relevant_retrieved),
            'num_relevant_total': len(relevant_doc_ids)
        }

    def evaluate_generation(
        self,
        query: str,
        generated_answer: str,
        context: str,
        reference_answer: str = None
    ) -> Dict:
        """
        Evaluate generation quality.

        Metrics:
        - Faithfulness: Are claims grounded in context?
        - Relevance: Does it answer the question?
        - Completeness: Does it cover key points?
        - Answer similarity: How close to reference? (if provided)

        Uses LLM-as-judge for automated evaluation.
        """
        results = {}

        # Faithfulness: Check if answer is grounded in context
        results['faithfulness'] = self._check_faithfulness(
            generated_answer, context
        )

        # Relevance: Does answer address the query?
        results['relevance'] = self._check_relevance(
            query, generated_answer
        )

        # If reference answer provided, compute similarity
        if reference_answer:
            results['answer_similarity'] = self._compute_answer_similarity(
                generated_answer, reference_answer
            )

        return results

    def _check_faithfulness(self, answer: str, context: str) -> Dict:
        """
        Check if answer is faithful to context (no hallucinations).

        Returns:
        - is_faithful: bool
        - unfaithful_claims: List of claims not in context
        - score: 0-1 faithfulness score
        """
        prompt = f"""Given the following context and answer, determine if the answer is faithful to the context.
An answer is faithful if every factual claim in the answer is supported by the context.

Context:
{context}

Answer:
{answer}

Evaluate faithfulness:
1. List each factual claim in the answer
2. For each claim, verify if it's supported by the context
3. Return your evaluation in JSON format:
{{
    "is_faithful": true/false,
    "claims": [
        {{"claim": "...", "supported": true/false, "evidence": "..."}},
        ...
    ],
    "score": 0.0-1.0,
    "explanation": "..."
}}

JSON Response:"""

        response = self.client.chat.completions.create(
            model="gpt-4-turbo-preview",
            temperature=0,
            response_format={"type": "json_object"},
            messages=[{"role": "user", "content": prompt}]
        )

        result = json.loads(response.choices[0].message.content)
        return result

    def _check_relevance(self, query: str, answer: str) -> Dict:
        """Check if answer relevantly addresses the query."""
        prompt = f"""Evaluate if the following answer relevantly addresses the query.

Query: {query}

Answer: {answer}

Return JSON evaluation:
{{
    "is_relevant": true/false,
    "relevance_score": 0.0-1.0,
    "explanation": "Why is/isn't this relevant?",
    "missing_aspects": ["What key aspects of the query are not addressed?"]
}}

JSON Response:"""

        response = self.client.chat.completions.create(
            model="gpt-4-turbo-preview",
            temperature=0,
            response_format={"type": "json_object"},
            messages=[{"role": "user", "content": prompt}]
        )

        result = json.loads(response.choices[0].message.content)
        return result

    def _compute_answer_similarity(
        self, generated: str, reference: str
    ) -> float:
        """
        Compute semantic similarity between generated and reference answer.
        Uses embedding-based similarity.
        """
        from openai import OpenAI
        client = OpenAI()

        # Get embeddings
        embeddings = client.embeddings.create(
            model="text-embedding-3-large",
            input=[generated, reference]
        )

        gen_emb = np.array(embeddings.data[0].embedding)
        ref_emb = np.array(embeddings.data[1].embedding)

        # Cosine similarity
        similarity = np.dot(gen_emb, ref_emb) / (
            np.linalg.norm(gen_emb) * np.linalg.norm(ref_emb)
        )

        return float(similarity)

    def evaluate_end_to_end(
        self,
        test_cases: List[Dict],
        rag_system
    ) -> Dict:
        """
        Run full evaluation on a test set.

        test_cases format:
        [
            {
                "query": "...",
                "relevant_doc_ids": [...],
                "reference_answer": "..."  # optional
            },
            ...
        ]

        Returns aggregated metrics across all test cases.
        """
        retrieval_results = []
        generation_results = []

        for case in test_cases:
            query = case['query']

            # Run RAG system
            result = rag_system.query(query)

            # Evaluate retrieval
            retrieval_metrics = self.evaluate_retrieval(
                query=query,
                retrieved_docs=result.get('retrieved_docs', []),
                relevant_doc_ids=set(case['relevant_doc_ids'])
            )
            retrieval_results.append(retrieval_metrics)

            # Evaluate generation
            generation_metrics = self.evaluate_generation(
                query=query,
                generated_answer=result['answer'],
                context=result.get('context', ''),
                reference_answer=case.get('reference_answer')
            )
            generation_results.append(generation_metrics)

        # Aggregate metrics
        return {
            'retrieval': self._aggregate_metrics(retrieval_results),
            'generation': self._aggregate_metrics(generation_results),
            'num_test_cases': len(test_cases)
        }

    def _aggregate_metrics(self, results: List[Dict]) -> Dict:
        """Compute mean and std for all numeric metrics."""
        aggregated = {}

        # Get all metric keys
        all_keys = set()
        for result in results:
            all_keys.update(result.keys())

        for key in all_keys:
            values = []
            for result in results:
                val = result.get(key)
                # Only aggregate numeric values
                if isinstance(val, (int, float)):
                    values.append(val)

            if values:
                aggregated[f'{key}_mean'] = np.mean(values)
                aggregated[f'{key}_std'] = np.std(values)

        return aggregated


# Usage Example
evaluator = RAGEvaluator()

# Golden dataset example
test_cases = [
    {
        "query": "What are the benefits of hybrid search?",
        "relevant_doc_ids": {"doc_42", "doc_103", "doc_87"},
        "reference_answer": "Hybrid search combines dense and sparse retrieval..."
    },
    # ... more test cases
]

# Evaluate your RAG system
results = evaluator.evaluate_end_to_end(
    test_cases=test_cases,
    rag_system=rag_system
)

print("Retrieval Metrics:")
print(f"  Precision@10: {results['retrieval']['precision@k_mean']:.2%} ± {results['retrieval']['precision@k_std']:.2%}")
print(f"  Recall@10: {results['retrieval']['recall@k_mean']:.2%}")
print(f"  MRR: {results['retrieval']['mrr_mean']:.3f}")
print(f"  NDCG: {results['retrieval']['ndcg_mean']:.3f}")

print("\nGeneration Metrics:")
print(f"  Faithfulness: {results['generation']['faithfulness_score_mean']:.2%}")
print(f"  Relevance: {results['generation']['relevance_score_mean']:.2%}")