Advanced RAG

"""Hybrid Search: BM25 + Vector Search with Reciprocal Rank Fusion"""
from rank_bm25 import BM25Okapi
import numpy as np
from openai import OpenAI
import chromadb

client = OpenAI()

class HybridSearcher:
    def __init__(self, documents: list[dict]):
        """
        documents: list of {"content": str, "metadata": dict}
        """
        self.documents = documents
        self.contents = [d["content"] for d in documents]
        
        # Build BM25 index (sparse/keyword search)
        tokenized = [doc.lower().split() for doc in self.contents]
        self.bm25 = BM25Okapi(tokenized)
        
        # Build vector index
        self.chroma = chromadb.Client()
        self.collection = self.chroma.create_collection("hybrid")
        
        # Embed and store
        embeddings = self._embed(self.contents)
        self.collection.add(
            ids=[f"doc_{i}" for i in range(len(documents))],
            documents=self.contents,
            embeddings=embeddings,
            metadatas=[d["metadata"] for d in documents]
        )
    
    def _embed(self, texts: list[str]) -> list[list[float]]:
        response = client.embeddings.create(
            input=texts, model="text-embedding-3-small"
        )
        return [d.embedding for d in response.data]
    
    def search(self, query: str, top_k: int = 10, alpha: float = 0.7) -> list[dict]:
        """
        Hybrid search with Reciprocal Rank Fusion (RRF).
        
        alpha: weight for vector search (1-alpha for BM25)
        """
        # --- Sparse search (BM25) ---
        bm25_scores = self.bm25.get_scores(query.lower().split())
        bm25_ranked = np.argsort(bm25_scores)[::-1][:top_k * 2]
        
        # --- Dense search (vector) ---
        vector_results = self.collection.query(
            query_texts=[query], n_results=top_k * 2
        )
        # Map IDs back to indices
        vector_ranked = [
            int(id.split("_")[1]) for id in vector_results["ids"][0]
        ]
        
        # --- Reciprocal Rank Fusion ---
        k = 60  # RRF constant
        scores = {}
        
        for rank, doc_idx in enumerate(bm25_ranked):
            scores[doc_idx] = scores.get(doc_idx, 0) + (1 - alpha) / (k + rank + 1)
        
        for rank, doc_idx in enumerate(vector_ranked):
            scores[doc_idx] = scores.get(doc_idx, 0) + alpha / (k + rank + 1)
        
        # Sort by combined score
        ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:top_k]
        
        return [
            {
                "content": self.contents[idx],
                "metadata": self.documents[idx]["metadata"],
                "score": score,
                "rank": i + 1
            }
            for i, (idx, score) in enumerate(ranked)
        ]

# Usage
searcher = HybridSearcher(documents)
results = searcher.search("authentication error 403", alpha=0.6)

"""Reranking with Cross-Encoders"""
from sentence_transformers import CrossEncoder

# Load a cross-encoder reranker
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

def retrieve_and_rerank(
    query: str, 
    initial_results: list[dict],  # From hybrid search
    top_k: int = 5
) -> list[dict]:
    """
    Two-stage retrieval:
    1. Fast retrieval (bi-encoder/BM25) → top 20-50 candidates
    2. Accurate reranking (cross-encoder) → top 5 final results
    """
    # Prepare pairs for cross-encoder
    pairs = [(query, r["content"]) for r in initial_results]
    
    # Score all pairs (cross-encoder sees query + doc together)
    scores = reranker.predict(pairs)
    
    # Attach scores and sort
    for result, score in zip(initial_results, scores):
        result["rerank_score"] = float(score)
    
    reranked = sorted(
        initial_results, 
        key=lambda x: x["rerank_score"], 
        reverse=True
    )
    
    return reranked[:top_k]

# --- Full pipeline ---
# Stage 1: Fast retrieval (get 20 candidates)
candidates = searcher.search(query="How to fix memory leaks?", top_k=20)

# Stage 2: Accurate reranking (narrow to top 5)
final_results = retrieve_and_rerank(
    query="How to fix memory leaks?",
    initial_results=candidates,
    top_k=5
)

for r in final_results:
    print(f"[{r['rerank_score']:.3f}] {r['content'][:80]}...")

"""Query Expansion and HyDE"""
from openai import OpenAI

client = OpenAI()

def expand_query(query: str, n_expansions: int = 3) -> list[str]:
    """Generate multiple query variations to improve recall."""
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{
            "role": "user",
            "content": (
                f"Generate {n_expansions} different ways to ask this question. "
                f"Include synonyms and related terms.\n\n"
                f"Original: {query}\n\n"
                f"Variations (one per line):"
            )
        }],
        temperature=0.7
    )
    
    variations = response.choices[0].message.content.strip().split("\n")
    return [query] + [v.strip("- 123.") for v in variations if v.strip()]


def hyde_search(query: str, collection, top_k: int = 5):
    """HyDE: Generate hypothetical answer, embed that instead."""
    # Generate hypothetical answer
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{
            "role": "user",
            "content": (
                f"Write a short, factual paragraph that would answer "
                f"this question:\n\n{query}\n\n"
                f"Write as if you're a technical documentation page."
            )
        }],
        temperature=0.3
    )
    
    hypothetical_answer = response.choices[0].message.content
    
    # Search using the hypothetical answer (closer to real docs)
    results = collection.query(
        query_texts=[hypothetical_answer],
        n_results=top_k
    )
    
    return results, hypothetical_answer


# --- Multi-query retrieval ---
def multi_query_retrieve(query: str, collection, top_k: int = 5):
    """Retrieve using multiple query variations, deduplicate."""
    queries = expand_query(query)
    
    all_results = {}  # doc_id → (content, best_score)
    
    for q in queries:
        results = collection.query(query_texts=[q], n_results=top_k)
        for doc_id, doc, dist in zip(
            results["ids"][0], 
            results["documents"][0],
            results["distances"][0]
        ):
            score = 1 - dist
            if doc_id not in all_results or score > all_results[doc_id][1]:
                all_results[doc_id] = (doc, score)
    
    # Sort by best score, return top_k
    ranked = sorted(all_results.items(), key=lambda x: x[1][1], reverse=True)
    return [{"id": k, "content": v[0], "score": v[1]} for k, v in ranked[:top_k]]

def compress_context(query: str, chunks: list[str]) -> list[str]:
    """Extract only query-relevant portions from retrieved chunks."""
    compressed = []
    
    for chunk in chunks:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{
                "role": "user",
                "content": (
                    f"Extract ONLY the parts of this text that are relevant "
                    f"to answering: '{query}'\n\n"
                    f"Text: {chunk}\n\n"
                    f"Relevant extract (or 'IRRELEVANT' if nothing matches):"
                )
            }],
            temperature=0,
            max_tokens=200
        )
        
        extract = response.choices[0].message.content.strip()
        if extract != "IRRELEVANT":
            compressed.append(extract)
    
    return compressed

# Reduces 5 chunks of 500 tokens each → maybe 3 chunks of 100 tokens
# More focused context = better generation quality

"""Parent-Child Chunking Strategy"""
from langchain.text_splitter import RecursiveCharacterTextSplitter

# Create parent chunks (large, for context)
parent_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1500, chunk_overlap=200
)

# Create child chunks (small, for retrieval)
child_splitter = RecursiveCharacterTextSplitter(
    chunk_size=300, chunk_overlap=50
)

def build_parent_child_index(documents: list[str], collection):
    """Index child chunks but store parent references."""
    all_children = []
    parent_store = {}  # parent_id → parent_text
    
    for doc_idx, doc in enumerate(documents):
        # Split into parents
        parents = parent_splitter.split_text(doc)
        
        for p_idx, parent in enumerate(parents):
            parent_id = f"doc{doc_idx}_parent{p_idx}"
            parent_store[parent_id] = parent
            
            # Split parent into children
            children = child_splitter.split_text(parent)
            
            for c_idx, child in enumerate(children):
                all_children.append({
                    "id": f"{parent_id}_child{c_idx}",
                    "content": child,
                    "parent_id": parent_id
                })
    
    # Index children (what we search against)
    collection.add(
        ids=[c["id"] for c in all_children],
        documents=[c["content"] for c in all_children],
        metadatas=[{"parent_id": c["parent_id"]} for c in all_children]
    )
    
    return parent_store

def search_with_parents(query: str, collection, parent_store, top_k: int = 3):
    """Search children, return parents (deduped)."""
    results = collection.query(query_texts=[query], n_results=top_k * 2)
    
    # Get unique parents
    seen_parents = set()
    parent_chunks = []
    
    for meta in results["metadatas"][0]:
        parent_id = meta["parent_id"]
        if parent_id not in seen_parents:
            seen_parents.add(parent_id)
            parent_chunks.append(parent_store[parent_id])
        
        if len(parent_chunks) >= top_k:
            break
    
    return parent_chunks