Building a RAG Pipeline with Gemma 4 and ChromaDB: Implementation

from langchain.document_loaders import TextLoader, PyPDFLoader
from langchain.schema import Document
from pathlib import Path
from typing import List
import logging
 
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
 
def load_documents(paths: List[str]) -> List[Document]:
    """
    Load documents from disk. Supports .txt, .md, and .pdf.
    
    Returns a flat list of LangChain Document objects.
    Missing or unreadable files are logged and skipped.
    """
    docs: List[Document] = []
 
    for path_str in paths:
        path = Path(path_str)
 
        if not path.exists():
            logger.warning(f"File not found, skipping: {path}")
            continue
 
        try:
            if path.suffix == ".pdf":
                loader = PyPDFLoader(str(path))
            else:
                loader = TextLoader(str(path), encoding="utf-8")
 
            loaded = loader.load()
            docs.extend(loaded)
            logger.info(f"Loaded {len(loaded)} page(s) from {path.name}")
 
        except Exception as e:
            logger.error(f"Failed to load {path}: {e}")
            continue
 
    logger.info(f"Total documents loaded: {len(docs)}")
    return docs

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document
from typing import List
 
def split_documents(
    docs: List[Document],
    chunk_size: int = 800,
    chunk_overlap: int = 100
) -> List[Document]:
    """
    Split documents into overlapping chunks.
    
    chunk_size: target characters per chunk (not tokens — adjust if
                your model's tokenizer differs significantly)
    chunk_overlap: characters shared between adjacent chunks,
                   preventing context loss at boundaries
    """
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", ". ", " ", ""],
    )
 
    chunks = splitter.split_documents(docs)
    logger.info(f"Split into {len(chunks)} chunks "
                f"(avg {sum(len(c.page_content) for c in chunks) // len(chunks)} chars)")
    return chunks

from langchain_community.vectorstores import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain.schema import Document
from typing import List, Optional, Tuple
import os
 
PERSIST_DIR = "./chroma_db"
EMBED_MODEL = "nomic-embed-text"  # Fast, good quality for RAG
 
def get_embeddings() -> OllamaEmbeddings:
    return OllamaEmbeddings(model=EMBED_MODEL)
 
def create_vector_store(chunks: List[Document]) -> Chroma:
    """
    Build ChromaDB from document chunks. Persists to disk automatically.
    This is slow on first run (embeds all chunks) — subsequent loads are instant.
    """
    if not chunks:
        raise ValueError("No chunks provided. Load and split documents first.")
 
    logger.info(f"Embedding {len(chunks)} chunks with {EMBED_MODEL}...")
    
    store = Chroma.from_documents(
        documents=chunks,
        embedding=get_embeddings(),
        persist_directory=PERSIST_DIR,
    )
 
    logger.info(f"Vector store created and persisted to {PERSIST_DIR}")
    return store
 
def load_existing_vector_store() -> Optional[Chroma]:
    """
    Load a previously built vector store from disk.
    Returns None if no store exists yet.
    """
    if not os.path.exists(PERSIST_DIR):
        logger.info("No existing vector store found.")
        return None
 
    store = Chroma(
        persist_directory=PERSIST_DIR,
        embedding_function=get_embeddings(),
    )
    logger.info("Loaded existing vector store from disk.")
    return store
 
def retrieve_relevant_chunks(
    store: Chroma,
    query: str,
    k: int = 4,
    score_threshold: float = 0.3
) -> List[Tuple[Document, float]]:
    """
    Retrieve the top-k most relevant chunks for a query.
    
    Chunks below score_threshold are filtered out to avoid
    injecting irrelevant context into the prompt.
    """
    results = store.similarity_search_with_relevance_scores(query, k=k)
    
    filtered = [(doc, score) for doc, score in results if score >= score_threshold]
    
    if not filtered:
        logger.warning(f"No chunks above threshold {score_threshold} for query: {query\!r}")
    else:
        logger.info(f"Retrieved {len(filtered)} chunks (scores: "
                    f"{[f'{s:.2f}' for _, s in filtered]})")
    
    return filtered

import ollama
from typing import List, Tuple
from langchain.schema import Document
 
GEMMA4_MODEL = "gemma4"
ANTIGRAVITY_MODEL = "gemma-2.5-pro-preview-05-06"  # Cloud fallback
 
def build_context(results: List[Tuple[Document, float]]) -> str:
    """Concatenate retrieved chunks into a single context block."""
    if not results:
        return ""
    
    parts = []
    for i, (doc, score) in enumerate(results, 1):
        source = doc.metadata.get("source", "unknown")
        parts.append(f"[Source {i} — {source} (relevance: {score:.2f})]\n{doc.page_content}")
    
    return "\n\n---\n\n".join(parts)
 
def generate_answer(
    query: str,
    context: str,
    stream: bool = False
) -> str:
    """
    Generate an answer using local Gemma 4 via Ollama.
    
    The prompt structure is tuned for factual, document-grounded answers.
    Gemma 4 responds well to explicit instructions not to fabricate.
    """
    if context:
        system_prompt = (
            "You are a helpful assistant with access to specific document excerpts. "
            "Answer based on those excerpts only. "
            "If the answer isn't in the excerpts, say so clearly — do not guess."
        )
        user_prompt = f"""Document excerpts:
{context}
 
---
 
Question: {query}
 
Answer based on the excerpts above."""
    else:
        system_prompt = "You are a helpful assistant."
        user_prompt = (
            f"The following question has no relevant document context available.\n\n"
            f"Question: {query}\n\n"
            f"Answer as best you can from general knowledge, and note that no documents were found."
        )
 
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]
 
    if stream:
        # Stream tokens to stdout for interactive use
        full_response = ""
        for chunk in ollama.chat(model=GEMMA4_MODEL, messages=messages, stream=True):
            token = chunk["message"]["content"]
            print(token, end="", flush=True)
            full_response += token
        print()  # newline after stream ends
        return full_response
    else:
        response = ollama.chat(model=GEMMA4_MODEL, messages=messages)
        return response["message"]["content"]
 
def generate_answer_cloud(query: str, context: str) -> str:
    """
    Cloud fallback using Google Antigravity API.
    Use when Gemma 4 struggles with complex multi-hop reasoning.
    """
    import google.generativeai as genai
    import os
 
    genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
    model = genai.GenerativeModel(ANTIGRAVITY_MODEL)
 
    if context:
        prompt = (
            f"Document excerpts:\n{context}\n\n"
            f"---\n\nQuestion: {query}\n\n"
            f"Answer based on the excerpts above."
        )
    else:
        prompt = query
 
    response = model.generate_content(prompt)
    return response.text

import os
from pathlib import Path
 
def run_rag_pipeline(
    document_paths: List[str],
    query: str,
    use_cloud: bool = False,
    stream: bool = True,
    force_rebuild: bool = False
) -> str:
    """
    Full RAG pipeline: load → chunk → retrieve → generate.
    
    document_paths: files to index (PDF, TXT, MD)
    query: the user's question
    use_cloud: use Antigravity API instead of local Gemma 4
    stream: stream tokens to stdout (local Gemma 4 only)
    force_rebuild: ignore cached vector store and rebuild from scratch
    """
    # --- Load or build vector store ---
    store = None
    if not force_rebuild:
        store = load_existing_vector_store()
 
    if store is None:
        docs = load_documents(document_paths)
        if not docs:
            return "Error: No documents could be loaded."
        
        chunks = split_documents(docs)
        store = create_vector_store(chunks)
 
    # --- Retrieve relevant chunks ---
    results = retrieve_relevant_chunks(store, query, k=4, score_threshold=0.3)
    context = build_context(results)
 
    # --- Generate answer ---
    if use_cloud:
        answer = generate_answer_cloud(query, context)
    else:
        answer = generate_answer(query, context, stream=stream)
 
    return answer
 
def main():
    """Interactive RAG session."""
    import sys
 
    # Default documents — modify as needed
    docs = [
        "./docs/manual.pdf",
        "./docs/faq.txt",
        "./docs/changelog.md",
    ]
 
    # Filter to existing files only
    docs = [d for d in docs if Path(d).exists()]
    if not docs:
        print("No documents found in ./docs/ — place your files there and retry.")
        sys.exit(1)
 
    print(f"RAG pipeline ready. Documents: {[Path(d).name for d in docs]}")
    print("Type your question (or 'quit' to exit, 'rebuild' to re-index documents):\n")
 
    while True:
        try:
            query = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nExiting.")
            break
 
        if not query:
            continue
        if query.lower() == "quit":
            break
        if query.lower() == "rebuild":
            run_rag_pipeline(docs, "warmup", force_rebuild=True, stream=False)
            print("Vector store rebuilt.\n")
            continue
 
        print("\nGemma 4: ", end="")
        answer = run_rag_pipeline(docs, query, stream=True)
        print()
 
if __name__ == "__main__":
    main()

ollama pull mxbai-embed-large