Building a RAG Pipeline with Antigravity— Unlock Your Company's Knowledge with Vector Search and LLMs

# semantic_chunker.py — Advanced chunking that preserves semantic coherence
 
from langchain.text_splitter import RecursiveCharacterTextSplitter  # used by the parent-child splitter below
from langchain_experimental.text_splitter import SemanticChunker
from langchain_google_genai import GoogleGenerativeAIEmbeddings
import os
 
def create_semantic_chunks(documents: list) -> list:
    """Embedding-based semantic chunking"""
    embeddings = GoogleGenerativeAIEmbeddings(
        model="gemini-embedding-001",
        google_api_key=os.getenv("GOOGLE_API_KEY"),
        output_dimensionality=1536,
    )
 
    # Split where semantic distance between sentences exceeds threshold
    chunker = SemanticChunker(
        embeddings,
        breakpoint_threshold_type="percentile",
        breakpoint_threshold_amount=90,  # Split at top 10% distance
    )
 
    chunks = []
    for doc in documents:
        split = chunker.create_documents(
            [doc.page_content],
            metadatas=[doc.metadata],
        )
        chunks.extend(split)
 
    return chunks
 
# --- Parent-Child Chunking Strategy ---
class ParentChildChunker:
    """
    Use large 'parent chunks' to preserve context and
    small 'child chunks' for precise retrieval — the best of both worlds.
    """
 
    def __init__(self, parent_size=2000, child_size=400, overlap=100):
        self.parent_splitter = RecursiveCharacterTextSplitter(
            chunk_size=parent_size,
            chunk_overlap=0,
        )
        self.child_splitter = RecursiveCharacterTextSplitter(
            chunk_size=child_size,
            chunk_overlap=overlap,
        )
 
    def split(self, documents: list) -> tuple:
        """Generate parent and child chunks simultaneously"""
        parent_chunks = []
        child_chunks = []
 
        for doc in documents:
            parents = self.parent_splitter.split_documents([doc])
            for idx, parent in enumerate(parents):
                parent.metadata["parent_id"] = f"{doc.metadata.get('source', 'unknown')}_{idx}"
                parent_chunks.append(parent)
 
                # Split parent into smaller child chunks
                children = self.child_splitter.split_documents([parent])
                for child in children:
                    child.metadata["parent_id"] = parent.metadata["parent_id"]
                    child_chunks.append(child)
 
        return parent_chunks, child_chunks
 
# Expected behavior:
# - Search is performed against child chunks for precision
# - The parent chunk of each hit is passed as context to the LLM
# → Achieves both retrieval precision and rich context

# hybrid_retriever.py — Hybrid search combining vector search and BM25
 
from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever
from langchain_community.vectorstores import Chroma
from langchain_google_genai import GoogleGenerativeAIEmbeddings
import os
 
def create_hybrid_retriever(chunks: list, k: int = 5):
    """Hybrid retriever combining vector search and BM25"""
 
    # Deliberately no task_type here. The library assigns
    # RETRIEVAL_DOCUMENT to embed_documents and RETRIEVAL_QUERY
    # to embed_query on its own — more on that later.
    embeddings = GoogleGenerativeAIEmbeddings(
        model="gemini-embedding-001",
        google_api_key=os.getenv("GOOGLE_API_KEY"),
        output_dimensionality=1536,
    )
 
    # 1. Vector search retriever
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        collection_name="hybrid_search",
    )
    vector_retriever = vectorstore.as_retriever(
        search_type="mmr",          # Maximal Marginal Relevance
        search_kwargs={"k": k, "fetch_k": k * 3},
    )
 
    # 2. BM25 keyword search retriever
    bm25_retriever = BM25Retriever.from_documents(
        chunks,
        k=k,
    )
 
    # 3. Ensemble (weighted combination)
    ensemble = EnsembleRetriever(
        retrievers=[vector_retriever, bm25_retriever],
        weights=[0.6, 0.4],  # Slightly favor vector search
    )
 
    return ensemble
 
# --- Reranking with Cohere Reranker ---
from langchain.retrievers import ContextualCompressionRetriever
from langchain_cohere import CohereRerank
 
def add_reranking(base_retriever, top_n: int = 3):
    """Re-score search results with a reranker"""
    reranker = CohereRerank(
        model="rerank-v3.5",
        cohere_api_key=os.getenv("COHERE_API_KEY"),
        top_n=top_n,
    )
 
    return ContextualCompressionRetriever(
        base_compressor=reranker,
        base_retriever=base_retriever,
    )
 
# Usage example:
# retriever = create_hybrid_retriever(chunks, k=10)
# reranked_retriever = add_reranking(retriever, top_n=3)
# results = reranked_retriever.invoke("How do I verify Stripe webhook signatures?")
#
# Expected behavior:
# → BM25 picks up candidates matching "Stripe", "webhook", "signatures"
# → Vector search also finds semantically similar documents
# → Reranker narrows down to the top 3 most relevant documents

# rag_chain.py — Building the RAG chain
 
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
 
# Keep the model name in one place — a deprecation notice then costs you one line
CHAT_MODEL = "gemini-2.5-pro"   # scheduled for retirement on or after October 16, 2026
 
def build_rag_chain(retriever):
    """Build a RAG chain that generates answers from retrieved context"""
 
    # --- Prompt template ---
    template = ChatPromptTemplate.from_messages([
        ("system", """You are an assistant that answers questions accurately based on internal knowledge.
 
Follow these rules strictly:
1. Base your answer ONLY on the provided context
2. If the context doesn't contain the answer, explicitly state "This information was not found in the knowledge base"
3. Never mix in speculation or external knowledge
4. Cite the source document for each claim
5. Include code examples for technical content
 
Context:
{context}"""),
        ("human", "{question}"),
    ])
 
    # --- LLM ---
    llm = ChatGoogleGenerativeAI(
        model=CHAT_MODEL,
        temperature=0.1,        # Low temperature for factual responses
        max_output_tokens=4096,
    )
 
    # --- Chain construction ---
    def format_docs(docs):
        """Format search results into a context string"""
        formatted = []
        for i, doc in enumerate(docs, 1):
            source = doc.metadata.get("source", "unknown")
            formatted.append(
                f"[Source {i}: {source}]\n{doc.page_content}"
            )
        return "\n\n---\n\n".join(formatted)
 
    chain = (
        {
            "context": retriever | format_docs,
            "question": RunnablePassthrough(),
        }
        | template
        | llm
        | StrOutputParser()
    )
 
    return chain
 
# --- Streaming version ---
def query_with_streaming(chain, question: str):
    """Display answer with streaming"""
    print(f"\n📝 Question: {question}\n")
    print("=" * 60)
    for chunk in chain.stream(question):
        print(chunk, end="", flush=True)
    print("\n" + "=" * 60)
 
# Usage example:
# chain = build_rag_chain(reranked_retriever)
# query_with_streaming(chain, "What's the setup procedure for a new project?")
#
# Expected output:
# 📝 Question: What's the setup procedure for a new project?
# ============================================================
# Based on our internal knowledge base, new project setup follows
# these steps:
# 1. Create a repository on GitHub...
# [Source: docs/setup-guide.md]
# ============================================================

# rag_evaluator.py — RAG pipeline evaluation framework
 
from dataclasses import dataclass
from langchain_google_genai import ChatGoogleGenerativeAI
from pydantic import BaseModel, Field
 
@dataclass
class EvalResult:
    question: str
    expected: str
    actual: str
    faithfulness: float    # Is it faithful to the context? (0-1)
    relevance: float       # Is it relevant to the question? (0-1)
 
class Verdict(BaseModel):
    """Schema the judge model is required to return"""
    score: float = Field(description="A score between 0.0 and 1.0")
    reason: str = Field(description="Why that score was given")
 
class RAGEvaluator:
    """Evaluate RAG answers using LLM-as-a-Judge pattern"""
 
    def __init__(self):
        # Calling json.loads on the raw response breaks the moment the model
        # wraps its JSON in a ```json fence. Structured output removes the
        # parsing step entirely.
        self.judge = ChatGoogleGenerativeAI(
            model="gemini-2.5-pro",
            temperature=0,
        ).with_structured_output(Verdict)
 
    def evaluate_faithfulness(
        self, context: str, answer: str
    ) -> float:
        """Evaluate whether the answer is faithful to the context"""
        prompt = f"""Evaluate whether the following answer is based solely on the provided context.
 
Context:
{context}
 
Answer:
{answer}
 
Return a score from 0.0 (completely unsupported) to 1.0 (fully faithful)."""
 
        return self.judge.invoke(prompt).score
 
    def evaluate_relevance(
        self, question: str, answer: str
    ) -> float:
        """Evaluate whether the answer addresses the question"""
        prompt = f"""Evaluate whether the following answer appropriately addresses the question.
 
Question: {question}
Answer: {answer}
 
Return a score from 0.0 (irrelevant) to 1.0 (perfectly on point)."""
 
        return self.judge.invoke(prompt).score
 
    def run_eval_suite(
        self, chain, retriever, test_cases: list[dict]
    ) -> list[EvalResult]:
        """Run evaluation across a test case suite"""
        results = []
        for case in test_cases:
            # Retrieve relevant documents
            docs = retriever.invoke(case["question"])
            context = "\n".join(d.page_content for d in docs)
 
            # Generate answer
            answer = chain.invoke(case["question"])
 
            # Evaluate
            faithfulness = self.evaluate_faithfulness(context, answer)
            relevance = self.evaluate_relevance(case["question"], answer)
 
            results.append(EvalResult(
                question=case["question"],
                expected=case.get("expected", ""),
                actual=answer,
                faithfulness=faithfulness,
                relevance=relevance,
            ))
 
        # Print summary
        avg_faith = sum(r.faithfulness for r in results) / len(results)
        avg_rel = sum(r.relevance for r in results) / len(results)
        print(f"\n📊 Evaluation Summary ({len(results)} cases)")
        print(f"  Faithfulness: {avg_faith:.2f}")
        print(f"  Relevance:    {avg_rel:.2f}")
 
        return results
 
# Test case example:
# test_cases = [
#     {"question": "What's the deploy procedure?", "expected": "Run wrangler deploy..."},
#     {"question": "How do I generate API keys?", "expected": "From the admin panel..."},
# ]
# evaluator = RAGEvaluator()
# results = evaluator.run_eval_suite(chain, retriever, test_cases)