Optimizing RAG: Hybrid Search and Re-ranking for Low-Latency LLMs

import faiss
import numpy as np
from elasticsearch import Elasticsearch
from sentence_transformers import SentenceTransformer

# --- Assume these are pre-populated ---
# Elasticsearch client
es_client = Elasticsearch("http://localhost:9200")
# Sentence Transformer model for embedding queries
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
# FAISS index loaded from disk
faiss_index = faiss.read_index("my_faiss_index.bin")
# A mapping from FAISS index position to our document ID
# e.g., {0: 'doc_abc', 1: 'doc_xyz', ...}
index_to_doc_id = { ... }

# --- Implementation ---

def search_bm25(query_text: str, k: int = 50) -> list[tuple[str, float]]:
    """Performs a BM25 search and returns (doc_id, score) tuples."""
    response = es_client.search(
        index="my_documents",
        body={
            "size": k,
            "query": {
                "match": {
                    "content": query_text
                }
            }
        }
    )
    return [(hit['_id'], hit['_score']) for hit in response['hits']['hits']]

def search_vector(query_text: str, k: int = 50) -> list[tuple[str, float]]:
    """Performs a vector search and returns (doc_id, score) tuples."""
    query_vector = embedding_model.encode([query_text])
    # FAISS returns distances (lower is better), so we'll convert to similarity later
    distances, indices = faiss_index.search(query_vector.astype('float32'), k)
    
    results = []
    for i, dist in zip(indices[0], distances[0]):
        if i != -1: # FAISS can return -1 for no result
            doc_id = index_to_doc_id[i]
            # Convert distance to a similarity score (0-1), 1 / (1 + dist)
            similarity = 1.0 / (1.0 + dist)
            results.append((doc_id, similarity))
    return results

def reciprocal_rank_fusion(results_lists: list[list[tuple[str, float]]], k: int = 60) -> dict[str, float]:
    """Performs RRF on a list of search result lists."""
    fused_scores = {}
    
    # Create a map from doc_id to its rank in each list
    doc_ranks = {}
    for results in results_lists:
        for rank, (doc_id, _) in enumerate(results):
            if doc_id not in doc_ranks:
                doc_ranks[doc_id] = []
            doc_ranks[doc_id].append(rank + 1) # rank is 1-based
            
    # Calculate RRF score for each document
    for doc_id, ranks in doc_ranks.items():
        rrf_score = 0.0
        for rank in ranks:
            rrf_score += 1.0 / (k + rank)
        fused_scores[doc_id] = rrf_score
        
    return fused_scores

# --- Putting it together ---
def hybrid_search(query: str, top_k_retrieval: int = 100):
    bm25_results = search_bm25(query, k=top_k_retrieval)
    vector_results = search_vector(query, k=top_k_retrieval)

    # RRF requires ranked lists, so we don't need the scores here, just the order
    bm25_ranked_list = [(doc_id, score) for doc_id, score in bm25_results]
    vector_ranked_list = [(doc_id, score) for doc_id, score in vector_results]

    fused_scores = reciprocal_rank_fusion([bm25_ranked_list, vector_ranked_list])

    # Sort by RRF score in descending order
    sorted_fused_results = sorted(fused_scores.items(), key=lambda item: item[1], reverse=True)
    
    return sorted_fused_results

# --- Example Usage ---
query = "What is the p99 latency for the user-service API?"
retrieved_docs = hybrid_search(query)
print(f"Retrieved {len(retrieved_docs)} documents after fusion.")
print("Top 5 results:", retrieved_docs[:5])

from sentence_transformers.cross_encoder import CrossEncoder

# Load a lightweight, pre-trained cross-encoder
# These models are trained on tasks like question-answering and are ideal for re-ranking
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

# Assume 'retrieved_docs' is the output from our hybrid_search function
# and we have a function to get the actual text content of a document
def get_document_content(doc_id: str) -> str:
    # In a real system, this would fetch from a database or document store
    # For this example, we'll use a dummy dictionary
    dummy_db = {
        'doc1': 'The user-service API has a p99 latency of 250ms.',
        'doc2': 'System performance metrics are tracked in Grafana.',
        'doc3': 'API latency is a key performance indicator.',
        'doc4': 'The billing-service API latency is currently 150ms.'
        # ... more docs
    }
    return dummy_db.get(doc_id, "")

def rerank_documents(query: str, doc_ids: list[str], top_n: int = 5) -> list[tuple[str, float]]:
    """Re-ranks a list of document IDs using a cross-encoder."""
    # Create pairs of [query, document_content] for the cross-encoder
    query_doc_pairs = [(query, get_document_content(doc_id)) for doc_id in doc_ids]
    
    if not query_doc_pairs:
        return []
        
    # Predict scores. The model will output a single relevance score for each pair.
    scores = cross_encoder.predict(query_doc_pairs)
    
    # Combine doc_ids with their new scores
    reranked_results = list(zip(doc_ids, scores))
    
    # Sort by the new score in descending order
    reranked_results.sort(key=lambda x: x[1], reverse=True)
    
    return reranked_results[:top_n]

# --- Example Usage with the previous stage's output ---
# Let's say hybrid_search returned a list of (doc_id, rrf_score) tuples
retrieved_doc_ids = [doc_id for doc_id, score in retrieved_docs]

# We only re-rank the top, say, 50 candidates from the retrieval stage
rerank_candidates = retrieved_doc_ids[:50]

final_docs_for_llm = rerank_documents(query, rerank_candidates, top_n=5)

print("Final documents to be passed to LLM context:", final_docs_for_llm)
# Expected output might be: [('doc1', 0.98), ('doc3', 0.75), ('doc2', 0.21), ...]
# Notice how the most direct answer ('doc1') is pushed to the top.

from fastapi import FastAPI
from pydantic import BaseModel
import openai

# Assume all previous functions (search_bm25, search_vector, hybrid_search, rerank_documents, etc.) are defined here.
# Also assume openai.api_key is configured.

app = FastAPI()

class RAGQuery(BaseModel):
    query: str
    user_id: str # For logging, personalization, etc.

class RAGResponse(BaseModel):
    answer: str
    retrieved_doc_ids: list[str]

def build_llm_prompt(query: str, context_docs: list[str]) -> str:
    """Builds the final prompt for the LLM."""
    context = "\n\n---\n\n".join(context_docs)
    prompt = f"""
    You are a helpful AI assistant. Answer the user's question based on the following context.
    If the context does not contain the answer, state that you don't have enough information.

    Context:
    {context}

    Question: {query}

    Answer:
    """
    return prompt

@app.post("/query", response_model=RAGResponse)
async def execute_rag_pipeline(rag_query: RAGQuery):
    # STAGE 1: Hybrid Retrieval
    # We retrieve a larger set of candidates (e.g., 50-100)
    retrieved_results = hybrid_search(rag_query.query, top_k_retrieval=75)
    retrieved_doc_ids = [doc_id for doc_id, score in retrieved_results]

    # STAGE 2: Cross-Encoder Re-ranking
    # We re-rank only the top candidates to keep latency down
    rerank_candidates = retrieved_doc_ids[:25]
    # We only need the top 3-5 documents for the final context
    final_reranked_docs = rerank_documents(rag_query.query, rerank_candidates, top_n=3)
    final_doc_ids = [doc_id for doc_id, score in final_reranked_docs]

    # STAGE 3: Context Generation and LLM Invocation
    context_contents = [get_document_content(doc_id) for doc_id in final_doc_ids]
    
    prompt = build_llm_prompt(rag_query.query, context_contents)

    # In a real app, you'd use async OpenAI client
    response = openai.ChatCompletion.create(
        model="gpt-4-turbo-preview",
        messages=[
            {"role": "system", "content": "You are an expert Q&A system."},
            {"role": "user", "content": prompt}
        ],
        temperature=0.1
    )

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

    return RAGResponse(answer=answer, retrieved_doc_ids=final_doc_ids)