Implementing Multi-Stage Reranking for High Precision Retrieval Augmented Generation on Google Cloud Platform

import os
from google.cloud import storage, bigquery
import textract   # for extracting text from PDFs/Word docs

PROJECT_ID = "my-gcp-project"
BUCKET_NAME = "rag-corpus-bucket"
BQ_DATASET = "rag_dataset"
BQ_TABLE = "documents"

storage_client = storage.Client()
bq_client = bigquery.Client()

def upload_to_gcs(local_path, gcs_path):
    bucket = storage_client.bucket(BUCKET_NAME)
    blob = bucket.blob(gcs_path)
    blob.upload_from_filename(local_path)

def extract_text(file_path):
    return textract.process(file_path).decode("utf-8")

def load_to_bigquery(rows):
    table_ref = f"{PROJECT_ID}.{BQ_DATASET}.{BQ_TABLE}"
    errors = bq_client.insert_rows_json(table_ref, rows)
    if errors:
        raise RuntimeError(f"BigQuery insert errors: {errors}")

def ingest_directory(root_dir):
    rows = []
    for dirpath, _, filenames in os.walk(root_dir):
        for fn in filenames:
            local_path = os.path.join(dirpath, fn)
            gcs_path = f"raw/{fn}"
            upload_to_gcs(local_path, gcs_path)

            text = extract_text(local_path)
            rows.append({
                "doc_id": fn,
                "gcs_uri": f"gs://{BUCKET_NAME}/{gcs_path}",
                "content": text,
                "source_type": os.path.splitext(fn)[1].lower()
            })
    load_to_bigquery(rows)

# Run once
# ingest_directory("/data/my_corpus")

import nltk
nltk.download('punkt')
from nltk.tokenize import sent_tokenize

MAX_TOKENS = 300  # Adjust based on target LLM context window

def chunk_text(doc_id, content):
    sentences = sent_tokenize(content)
    chunks = []
    current = []
    current_len = 0

    for sent in sentences:
        token_len = len(sent.split())
        if current_len + token_len > MAX_TOKENS:
            chunks.append({
                "chunk_id": f"{doc_id}_c{len(chunks)}",
                "doc_id": doc_id,
                "content": " ".join(current)
            })
            current = [sent]
            current_len = token_len
        else:
            current.append(sent)
            current_len += token_len

    # Add the tail
    if current:
        chunks.append({
            "chunk_id": f"{doc_id}_c{len(chunks)}",
            "doc_id": doc_id,
            "content": " ".join(current)
        })
    return chunks

from vertexai.preview.language_models import TextEmbeddingModel
import numpy as np

model = TextEmbeddingModel.from_pretrained("textembedding-gecko@001")

def embed_chunks(chunks):
    texts = [c["content"] for c in chunks]
    embeddings = model.get_embeddings(texts)  # Returns list of float vectors
    for c, emb in zip(chunks, embeddings):
        c["embedding"] = emb
    return chunks

# Example for a single document
# doc_chunks = chunk_text("doc123", raw_text)
# embed_chunks(doc_chunks)

from google.cloud import aiplatform_v1beta1 as aiplatform

client = aiplatform.MatchServiceClient()
parent = f"projects/{PROJECT_ID}/locations/us-central1"

# 1. Create index
index = {
    "display_name": "rag-chunk-index",
    "metadata_schema_uri": "gs://google-cloud-aiplatform/schema/matchingengine/metadata_schema_v1.yaml",
    "metadata": {
        "algorithm_config": {"tree_ah_config": {"leaf_node_embedding_count": 1000}}
    },
}
operation = client.create_index(parent=parent, index=index)
index_name = operation.result().name

# 2. Upload embeddings (batch)
# Export embeddings from BigQuery to GCS as .jsonl, then use the import API

gcloud alpha discoveryengine data-stores create rag-search \
    --project=$PROJECT_ID \
    --location=global \
    --collection-id=default_collection

from google.cloud import discoveryengine_v1alpha as discoveryengine

client = discoveryengine.DocumentServiceClient()
parent = client.branch_path(PROJECT_ID, "global", "rag-search", "default_branch")

def upload_chunk(chunk):
    doc = discoveryengine.Document(
        name=client.document_path(PROJECT_ID, "global", "rag-search", "default_branch", chunk["chunk_id"]),
        schema_id="document",
        title=f"Chunk from {chunk['doc_id']}",
        content=chunk["content"],
    )
    client.create_document(parent=parent, document=doc)

# Iterate over all chunks
# for chunk in all_chunks:
#     upload_chunk(chunk)

def embed_query(query: str) -> list[float]:
    return model.get_embeddings([query])[0]

def dense_search(query_emb, top_k=50):
    request = {
        "index_endpoint": f"projects/{PROJECT_ID}/locations/us-central1/indexEndpoints/{INDEX_ENDPOINT_ID}",
        "deployed_index_id": DEPLOYED_INDEX_ID,
        "queries": [{"embedding": query_emb}],
        "neighbor_count": top_k,
    }
    response = client.find_neighbors(request=request)
    return [(neighbor.id, neighbor.distance) for neighbor in response.nearest_neighbors[0].neighbors]

def lexical_search(query: str, top_k=50):
    search_service = discoveryengine.SearchServiceClient()
    request = discoveryengine.SearchRequest(
        serving_config=f"projects/{PROJECT_ID}/locations/global/collections/default_collection/dataStores/rag-search/servingConfigs/default_config",
        query=query,
        page_size=top_k,
    )
    response = search_service.search(request)
    return [(hit.id, hit.relevance_score) for hit in response.results]

def merge_candidates(dense_hits, lexical_hits, max_candidates=100):
    # Simple union with score normalization
    merged = {}
    for doc_id, score in dense_hits + lexical_hits:
        if doc_id not in merged or merged[doc_id] > score:
            merged[doc_id] = score
    # Sort by normalized score
    sorted_hits = sorted(merged.items(), key=lambda x: x[1])
    return sorted_hits[:max_candidates]

FROM python:3.11-slim
WORKDIR /app
RUN pip install --no-cache-dir torch sentence-transformers fastapi uvicorn
COPY app.py .
EXPOSE 8080
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8080"]

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from sentence_transformers import CrossEncoder

app = FastAPI()
model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

class RerankRequest(BaseModel):
    query: str
    candidates: list[str]   # list of chunk IDs
    contents: dict[str, str]  # mapping id → text

class RerankResponse(BaseModel):
    scores: dict[str, float]

@app.post("/rerank", response_model=RerankResponse)
def rerank(req: RerankRequest):
    pairs = [(req.query, req.contents[cid]) for cid in req.candidates]
    scores = model.predict(pairs)   # higher = more relevant
    return RerankResponse(scores=dict(zip(req.candidates, scores)))

gcloud run deploy cross-encoder-reranker \
  --image=gcr.io/$PROJECT_ID/cross-encoder-reranker \
  --platform=managed \
  --region=us-central1 \
  --allow-unauthenticated

import requests

def stage1_rerank(query, candidate_ids, id_to_content):
    payload = {
        "query": query,
        "candidates": candidate_ids,
        "contents": id_to_content,
    }
    resp = requests.post("https://cross-encoder-reranker-xxxxx.run.app/rerank", json=payload)
    resp.raise_for_status()
    return resp.json()["scores"]

from vertexai.preview import generative_models

model = generative_models.GenerativeModel("gemini-1.5-flash-001")
def llm_score(query, passage):
    prompt = f"""You are a relevance judge for a knowledge‑base retrieval system.
Given the user question and a candidate passage, output a single number between 0 and 1 indicating how well the passage answers the question. Use the following format:

Question: <question>
Passage: <passage>
Score: <number>

Do not add any extra text.

Question: {query}
Passage: {passage}
Score:"""
    response = model.generate_content(prompt, temperature=0.0)
    # Extract the numeric token
    import re
    match = re.search(r"Score:\s*([\d\.]+)", response.text)
    return float(match.group(1)) if match else 0.0

from concurrent.futures import ThreadPoolExecutor

def stage2_rerank(query, top_ids, id_to_content):
    scores = {}
    with ThreadPoolExecutor(max_workers=8) as executor:
        futures = {
            executor.submit(llm_score, query, id_to_content[cid]): cid for cid in top_ids
        }
        for fut in futures:
            cid = futures[fut]
            try:
                scores[cid] = fut.result()
            except Exception as e:
                scores[cid] = 0.0
    # Sort descending
    return sorted(scores.items(), key=lambda kv: kv[1], reverse=True)

def build_prompt(query, passages):
    context = "\n---\n".join([f"[{i+1}] {p}" for i, p in enumerate(passages)])
    prompt = f"""You are an expert assistant. Use the following retrieved documents to answer the user's question. Cite each fact with the corresponding document number in brackets.

Context:
{context}

Question: {query}

Answer (include citations, e.g., [1]):"""
    return prompt

def generate_answer(prompt):
    model = generative_models.GenerativeModel("gemini-1.5-pro-001")
    response = model.generate_content(
        prompt,
        temperature=0.2,   # Low temp for factual answers
        max_output_tokens=1024,
        top_p=0.95,
    )
    return response.text.strip()

def rag_answer(query):
    # 1. Embed query
    q_emb = embed_query(query)

    # 2. Retrieve candidates
    dense_hits = dense_search(q_emb, top_k=50)
    lexical_hits = lexical_search(query, top_k=50)
    candidates = merge_candidates(dense_hits, lexical_hits, max_candidates=100)

    candidate_ids = [cid for cid, _ in candidates]
    # Fetch chunk contents from BigQuery (or cache)
    id_to_content = fetch_contents(candidate_ids)  # implement as a SELECT

    # 3. Stage‑1 rerank
    stage1_scores = stage1_rerank(query, candidate_ids, id_to_content)
    top_n1 = sorted(stage1_scores, key=stage1_scores.get, reverse=True)[:30]

    # 4. Stage‑2 LLM scorer
    top_n2 = stage2_rerank(query, top_n1, id_to_content)[:8]  # keep 8 passages

    top_passages = [id_to_content[cid] for cid, _ in top_n2]

    # 5. Build prompt + generate
    prompt = build_prompt(query, top_passages)
    answer = generate_answer(prompt)
    return answer