RAG Pipeline Workflow Fresh 🌱

Build a complete Retrieval-Augmented Generation pipeline using Gemini Embedding 2 for multimodal context retrieval.

Pipeline Architecture

Detailed Flow

Implementation

Step 1: Document Ingestion

from google import genai
from google.genai import types
import chromadb

# Initialize clients
embed_client = genai.Client()
chroma = chromadb.PersistentClient(path='./vectordb')
collection = chroma.get_or_create_collection(
    name='documents',
    metadata={'hnsw:space': 'cosine'}
)

def chunk_text(text, chunk_size=500, overlap=50):
    """Split text into overlapping chunks."""
    words = text.split()
    chunks = []
    for i in range(0, len(words), chunk_size - overlap):
        chunk = ' '.join(words[i:i + chunk_size])
        if chunk:
            chunks.append(chunk)
    return chunks

def ingest_document(doc_id, text):
    """Chunk, embed, and store a document."""
    chunks = chunk_text(text)

    for i, chunk in enumerate(chunks):
        result = embed_client.models.embed_content(
            model='gemini-embedding-2-preview',
            contents=chunk,
            config=types.EmbedContentConfig(
                task_type='RETRIEVAL_DOCUMENT',
                output_dimensionality=768  # Balance quality/cost
            )
        )

        collection.add(
            ids=[f'{doc_id}_chunk_{i}'],
            embeddings=[result.embeddings[0].values],
            documents=[chunk],
            metadatas=[{'doc_id': doc_id, 'chunk_index': i}]
        )

# Ingest your documents
ingest_document('doc1', 'Your long document text here...')

Step 2: Query and Retrieve

def retrieve(query, top_k=5):
    """Embed query and retrieve relevant chunks."""
    result = embed_client.models.embed_content(
        model='gemini-embedding-2-preview',
        contents=query,
        config=types.EmbedContentConfig(
            task_type='RETRIEVAL_QUERY',
            output_dimensionality=768
        )
    )

    results = collection.query(
        query_embeddings=[result.embeddings[0].values],
        n_results=top_k
    )

    return results['documents'][0]  # List of relevant chunks

Step 3: Generate Response

import google.generativeai as genai_llm

def rag_query(question):
    """Full RAG pipeline: retrieve context and generate answer."""
    # Retrieve relevant chunks
    context_chunks = retrieve(question, top_k=5)
    context = '\n\n'.join(context_chunks)

    # Generate answer with context
    model = genai_llm.GenerativeModel('gemini-2.0-flash')
    prompt = f"""Answer the question based on the provided context.
If the answer is not in the context, say so.

Context:
{context}

Question: {question}

Answer:"""

    response = model.generate_content(prompt)
    return response.text

# Usage
answer = rag_query('What are the key features of the product?')
print(answer)

Step 4: Multimodal RAG (Images + Text)

from google import genai
from google.genai import types

client = genai.Client()

def ingest_image_with_caption(doc_id, image_path, caption):
    """Embed an image with its caption as an aggregated vector."""
    with open(image_path, 'rb') as f:
        image_bytes = f.read()

    result = client.models.embed_content(
        model='gemini-embedding-2-preview',
        contents=[
            types.Content(parts=[
                types.Part.from_text(caption),
                types.Part.from_bytes(data=image_bytes, mime_type='image/jpeg'),
            ])
        ]
    )

    collection.add(
        ids=[doc_id],
        embeddings=[result.embeddings[0].values],
        documents=[caption],
        metadatas=[{'type': 'image', 'path': image_path}]
    )

Dimension Selection for RAG

Dimension	MTEB Score	Storage per 1M docs	Best For
3072	Top	~12 GB	Maximum accuracy
1536	68.17	~6 GB	Production default
768	67.99	~3 GB	Cost-effective RAG
256	66.19	~1 GB	Large-scale, budget

Task Type Pairing

Always use RETRIEVAL_DOCUMENT when indexing and RETRIEVAL_QUERY when searching. This asymmetric pairing is optimized for search quality.

Verification Checklist

• [ ] Documents chunked with appropriate overlap
• [ ] Embeddings generated with RETRIEVAL_DOCUMENT task type
• [ ] Vector database populated and queryable
• [ ] Queries embedded with RETRIEVAL_QUERY task type
• [ ] Top-K retrieval returning relevant results
• [ ] LLM generating accurate answers from context
• [ ] Dimensions consistent between ingestion and query

See Also

• Text Embeddings — Text embedding details
• Multimodal Embeddings — Cross-modal input
• Semantic Search — Pure search workflow
• Integrations — LangChain, LlamaIndex support