Generative AI Google Brain

gemini-rag-engine-rag-engine-weaviate

3 min read

Vertex AI RAG Engine with Weaviate

Source notebook

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

Creates a Vertex AI RAG Engine corpus backed by Weaviate and queries it with Gemini.

Summary

This notebook teaches how to configure Vertex AI RAG Engine with Weaviate as the vector database. It initializes the Vertex AI SDK, creates a RAG corpus using text-embedding-005, uploads or imports files from local text, Cloud Storage, and Google Drive, then retrieves context for generation. It demonstrates both Gemini GenerateContent with a RAG retrieval tool and a separate retrieval_query flow for passing contexts to other generation APIs.

Key code patterns

Initialize Vertex AI

PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
vertexai.init(project=PROJECT_ID, location=LOCATION)

Sets project and region before using Vertex AI RAG Engine.

Create Weaviate RAG corpus

embedding_model_config = rag.EmbeddingModelConfig(
    publisher_model="publishers/google/models/text-embedding-005"
)
vector_db = rag.Weaviate(
    weaviate_http_endpoint=WEAVIATE_HTTP_ENDPOINT,
    collection_name=COLLECTION_NAME,
    api_key=API_KEY,
)
rag_corpus = rag.create_corpus(
    display_name=DISPLAY_NAME,
    embedding_model_config=embedding_model_config,
    vector_db=vector_db,
)

Connects RAG Engine to a Weaviate collection using a first-party embedding model.

Upload and import files

rag.upload_file(
    corpus_name=rag_corpus.name,
    path="test.txt",
    display_name="test.txt",
    description="my test",
)
rag.import_files(
    corpus_name=rag_corpus.name,
    paths=[GCS_BUCKET],
    chunk_size=512,
    chunk_overlap=50,
)

Shows local upload and batch import with explicit chunking parameters.

Gemini with RAG tool

rag_resource = rag.RagResource(rag_corpus=rag_corpus.name)
rag_retrieval_tool = Tool.from_retrieval(
    retrieval=rag.Retrieval(
        source=rag.VertexRagStore(
            rag_resources=[rag_resource],
            similarity_top_k=10,
            vector_distance_threshold=0.4,
        )
    )
)
rag_model = GenerativeModel("gemini-2.0-flash", tools=[rag_retrieval_tool])
response = rag_model.generate_content(GENERATE_CONTENT_PROMPT)

Lets Gemini use retrieved RAG contexts when they pass the distance threshold.

Standalone retrieval

response = rag.retrieval_query(
    rag_resources=[rag_resource],
    text=RETRIEVAL_QUERY,
    similarity_top_k=10,
)
retrieved_context = " ".join(
    [context.text for context in response.contexts.contexts]
).replace("\n", "")

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

Models & APIs used

  • Models: publishers/google/models/text-embedding-005, gemini-2.0-flash
  • APIs / services: Vertex AI, Vertex AI RAG Engine API, Gemini GenerateContent API, Cloud Storage, Google Drive
  • SDKs / libraries: google-cloud-aiplatform, vertexai

When to use this

Use this pattern when you need Vertex AI RAG Engine to store embeddings in Weaviate and provide retrieved context to Gemini or another generation API.

Gotchas & caveats

  • The notebook requires an existing Google Cloud project with the Vertex AI API enabled.
  • After installing google-cloud-aiplatform, the notebook restarts the kernel before continuing.
  • Colab users must run auth.authenticate_user().
  • WEAVIATE_HTTP_ENDPOINT, COLLECTION_NAME, and API_KEY must be provided for the Weaviate instance.
  • The API_KEY field is described as a Secret Manager resource name placeholder.
  • Cloud Storage buckets and Google Drive files or folders need Viewer access for the Vertex RAG Data Service Agent.
  • rag.VertexRagStore and rag.retrieval_query comments state that currently only one corpus is allowed.
  • Imported files may take a few seconds to process before list_files shows them.

Best practices

  • Use GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_REGION environment variables as fallbacks for project and region.
  • Configure a Google first-party embedding model for the corpus.
  • Set chunk_size and chunk_overlap when importing files.
  • Check created resources with rag.get_corpus and rag.list_files.
  • Use similarity_top_k and vector_distance_threshold to control retrieved context.
  • Retrieve file IDs with rag.list_files when targeting specific rag_file_ids.
  • Guard cleanup with a delete_rag_corpus boolean before calling rag.delete_corpus.

Graph View

Backlinks