Vertex AI RAG Engine with Pinecone

Source notebook

Repo path: gemini/rag-engine/rag_engine_pinecone.ipynb · Open on GitHub · intermediate

Shows how to use Vertex AI RAG Engine with Pinecone as the vector database for Gemini retrieval.

Summary

This notebook demonstrates configuring Vertex AI RAG Engine to use a Pinecone index as its vector database. It walks through optional Pinecone and Secret Manager setup, creating or updating a RAG corpus, importing local, Cloud Storage, and Google Drive files, then querying Gemini with a Vertex RAG retrieval tool. It also shows how to run a standalone retrieval query and pass retrieved context to another generation API.

Key code patterns

Initialize Vertex AI and GenAI client

vertexai.init(project=PROJECT_ID, location=LOCATION)
client = genai.Client(
    vertexai=True,
    project=PROJECT_ID,
    location=LOCATION,
)

Sets the Google Cloud project and region for Vertex AI RAG Engine and Gemini calls.

Create Pinecone-backed RAG corpus

vector_db = rag.Pinecone(
    index_name=PINECONE_INDEX_NAME,
    api_key=secret_version.name,
)
rag_corpus = rag.create_corpus(
    display_name=DISPLAY_NAME,
    backend_config=rag.RagVectorDbConfig(vector_db=vector_db),
)

Connects a Vertex AI RAG corpus to a Pinecone index using a Secret Manager secret version.

Import files with chunking

response = rag.import_files(
    corpus_name=rag_corpus.name,
    paths=[GCS_BUCKET],
    transformation_config=rag.TransformationConfig(
        chunking_config=rag.ChunkingConfig(
            chunk_size=512,
            chunk_overlap=50,
        )
    ),
)

Shows document ingestion into the corpus with explicit chunk size and overlap settings.

Generate with RAG retrieval tool

rag_retrieval_tool = Tool(
    retrieval=Retrieval(
        vertex_rag_store=VertexRagStore(
            rag_resources=[VertexRagStoreRagResource(rag_corpus=rag_corpus.name)],
            similarity_top_k=10,
            vector_distance_threshold=0.4,
        )
    )
)
response = client.models.generate_content(
    model=MODEL_ID,
    contents=GENERATE_CONTENT_PROMPT,
    config=GenerateContentConfig(tools=[rag_retrieval_tool]),
)

Adds retrieved corpus context to Gemini generation with top-k and distance threshold controls.

Standalone retrieval query

response = rag.retrieval_query(
    rag_resources=[rag_resource],
    text=RETRIEVAL_QUERY,
    rag_retrieval_config=rag.RagRetrievalConfig(
        top_k=10,
        filter=rag.Filter(vector_distance_threshold=0.5),
    ),
)

Retrieves contexts separately so they can be passed to any SDK or model generation API.

Models & APIs used

  • Models: gemini-2.5-flash
  • APIs / services: Vertex AI, Secret Manager, Cloud Storage, Google Drive, Pinecone
  • SDKs / libraries: google-cloud-aiplatform, google-cloud-secret-manager, pinecone[grpc], google-genai, vertexai

When to use this

Use this pattern when you need Vertex AI RAG Engine to retrieve from a Pinecone vector index and ground Gemini responses on imported files.

Gotchas & caveats

  • Vertex AI API must be enabled for the Google Cloud project.
  • The notebook installs packages and requires a kernel restart before continuing.
  • Colab users must authenticate with google.colab.auth.authenticate_user().
  • A Pinecone API key is required, and the notebook stores it in Secret Manager for RAG Engine access.
  • For the first RAG corpus in a project, the service account may need to be provisioned by creating a corpus with an empty Pinecone config before granting secret access.
  • The RAG Engine service account must have roles/secretmanager.secretAccessor on the Pinecone API key secret.
  • For Cloud Storage and Google Drive imports, the Vertex RAG Data Service Agent must have Viewer access.
  • Pinecone index dimension must match the embedding model dimension chosen for the RAG corpus.
  • The notebook comments that currently only one corpus is allowed in rag_resources.

Best practices

  • Use Secret Manager instead of embedding the Pinecone API key directly in the RAG corpus config.
  • Grant the RAG Engine service account only the secretAccessor role needed to read the Pinecone API key.
  • Set chunk_size and chunk_overlap explicitly when importing files.
  • Use similarity_top_k and vector_distance_threshold to control retrieved context quality.
  • Clean up the RAG corpus with rag.delete_corpus when the created resources are no longer needed.