Generative AI Google Brain

gemini-grounding-grounding-with-vais

3 min read

Grounding with Vertex AI Search

Source notebook

Repo path: gemini/grounding/grounding_with_vais.ipynb · Open on GitHub · intermediate

Creates a Vertex AI Search engine and uses it to ground a Gemini response with retrieved context.

Summary

This notebook teaches how to create a Vertex AI Search data store from Cloud Storage documents, build a search engine, and use that engine as a retrieval tool for a Gemini model call. It demonstrates setup, authentication, data ingestion, engine creation, grounded generation, grounding metadata inspection, and cleanup of created resources.

Key code patterns

Initialize Vertex AI Search clients

client_options = ClientOptions(api_endpoint=f"{VAIS_LOCATION}-discoveryengine.googleapis.com") if VAIS_LOCATION != "global" else None
data_store_service_client = vais.DataStoreServiceClient(client_options=client_options)
document_service_client = vais.DocumentServiceClient(client_options=client_options)
engine_client = vais.EngineServiceClient(client_options=client_options)

Creates Discovery Engine service clients, using a regional endpoint when the data store is not global.

Create data store

data_store = vais.DataStore(
    display_name="Data Store for Vertex LLM Grounding demo",
    industry_vertical="GENERIC",
    solution_types=["SOLUTION_TYPE_SEARCH"],
    content_config="CONTENT_REQUIRED",
)
request = vais.CreateDataStoreRequest(
    parent=f"projects/{PROJECT_ID}/locations/{VAIS_LOCATION}/collections/default_collection",
    data_store=data_store,
    data_store_id=DATA_STORE_ID,
)
created_data_store = data_store_service_client.create_data_store(request).result()

Defines an unstructured generic search data store used as the corpus for grounding.

Import documents from GCS

branch_path = document_service_client.branch_path(
    project=PROJECT_ID, location=VAIS_LOCATION,
    data_store=DATA_STORE_ID, branch="default_branch")
document_service_client.import_documents(
    request=vais.ImportDocumentsRequest(
        parent=branch_path,
        gcs_source=vais.GcsSource(input_uris=[f"{GCS_SOURCE}/*"], data_schema="content"),
        reconciliation_mode=vais.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
    )
)

Ingests Cloud Storage files into the Vertex AI Search data store for later retrieval.

Create enterprise search engine

engine = vais.Engine(
    display_name="Engine for Vertex LLM Grounding demo",
    solution_type=vais.SolutionType.SOLUTION_TYPE_SEARCH,
    search_engine_config=vais.Engine.SearchEngineConfig(
        search_tier=vais.SearchTier.SEARCH_TIER_ENTERPRISE,
        search_add_ons=[vais.SearchAddOn.SEARCH_ADD_ON_LLM]),
    data_store_ids=[DATA_STORE_ID],
)
operation = engine_client.create_engine(vais.CreateEngineRequest(parent=parent, engine=engine, engine_id=engine_id))

Creates a search engine with enterprise tier and LLM add-on for grounding quality and extractive answers.

client = genai.Client(enterprise=True, project=PROJECT_ID, location=LOCATION)
vais_tool = Tool(retrieval=Retrieval(vertex_ai_search=VertexAISearch(
    engine=f"projects/{PROJECT_ID}/locations/global/collections/default_collection/engines/{engine_id}")))
response = client.models.generate_content(
    model=MODEL_ID,
    contents=PROMPT,
    config=GenerateContentConfig(tools=[vais_tool]),
)

Passes the Vertex AI Search engine as a retrieval tool so the model can ground its response in indexed content.

Inspect grounding metadata

for s in response.candidates[0].grounding_metadata.grounding_supports:
    display(Markdown(f"{s.segment.text} {s.grounding_chunk_indices}"))
for i, chunk in enumerate(response.candidates[0].grounding_metadata.grounding_chunks):
    display(Markdown(chunk.retrieved_context.text))
    print(chunk.retrieved_context.uri)

Shows supporting claim segments and retrieved contexts returned with the grounded model response.

Models & APIs used

  • Models: gemini-3.5-flash
  • APIs / services: Vertex AI, Vertex AI Search, Cloud Storage
  • SDKs / libraries: google-cloud-aiplatform, google-cloud-discoveryengine, google-genai

When to use this

Use this pattern when Gemini responses need to be grounded in a private or curated document corpus indexed by Vertex AI Search.

Gotchas & caveats

  • Vertex AI API must be enabled for the Google Cloud project.
  • Colab requires auth.authenticate_user(project_id=PROJECT_ID).
  • The notebook installs packages and restarts the runtime before continuing.
  • Vertex AI Search data stores support us, eu, and global regions in this notebook.
  • Enterprise search tier is required for extractive answers and advanced LLM features in this workflow.
  • Engine indexing can take a few minutes before it is ready for search or grounding.
  • The model may not always output grounding support even when retrieval succeeds.
  • The notebook creates billable cloud resources and deletes the engine and data store during cleanup.

Best practices

  • Use environment variables for project and region defaults when parameters are not provided.
  • Create the Vertex AI Search engine with SEARCH_TIER_ENTERPRISE and SEARCH_ADD_ON_LLM for grounding quality.
  • Verify the search engine is ready by sending a SearchRequest before using it for grounded generation.
  • Inspect grounding_supports and grounding_chunks to understand which retrieved documents support the answer.
  • Delete the created engine and data store after the tutorial.

Graph View

Backlinks