Create a Vertex AI Datastore and Search Engine

Source notebook

Repo path: search/create_datastore_and_search.ipynb · Open on GitHub · intermediate

Creates a Vertex AI Search datastore, imports GCS PDFs, creates an enterprise search engine, and queries it.

Summary

This notebook teaches how to create and populate a Vertex AI Search datastore, connect it to a search engine, and submit queries. The workflow imports Alphabet investor PDFs from Cloud Storage, creates an enterprise search engine with the LLM add-on, and runs a query with snippets, summaries, citations, query expansion, and spell correction.

Key code patterns

Regional Discovery Engine client

client_options = (
    ClientOptions(api_endpoint=f"{location}-discoveryengine.googleapis.com")
    if location != "global"
    else None
)
client = discoveryengine.DataStoreServiceClient(client_options=client_options)

Uses the location-specific Discovery Engine endpoint when the datastore is not global.

Create content datastore

data_store = discoveryengine.DataStore(
    display_name=data_store_name,
    industry_vertical=discoveryengine.IndustryVertical.GENERIC,
    content_config=discoveryengine.DataStore.ContentConfig.CONTENT_REQUIRED,
)
operation = client.create_data_store(
    request=discoveryengine.CreateDataStoreRequest(
        parent=client.collection_path(project_id, location, "default_collection"),
        data_store=data_store,
        data_store_id=data_store_id,
    )
)

Creates a generic datastore configured for document content.

Import GCS documents

parent = client.branch_path(project_id, location, data_store_id, "default_branch")
request = discoveryengine.ImportDocumentsRequest(
    parent=parent,
    gcs_source=discoveryengine.GcsSource(
        input_uris=[f"{gcs_uri}/*"], data_schema="content"
    ),
    reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
)
operation = client.import_documents(request=request)

Loads documents from a Cloud Storage folder into the datastore incrementally.

Create enterprise search engine

engine = discoveryengine.Engine(
    display_name=engine_name,
    solution_type=discoveryengine.SolutionType.SOLUTION_TYPE_SEARCH,
    industry_vertical=discoveryengine.IndustryVertical.GENERIC,
    data_store_ids=[data_store_id],
    search_engine_config=discoveryengine.Engine.SearchEngineConfig(
        search_tier=discoveryengine.SearchTier.SEARCH_TIER_ENTERPRISE,
        search_add_ons=[discoveryengine.SearchAddOn.SEARCH_ADD_ON_LLM],
    ),
)

Enables enterprise search tier and LLM add-on for summaries and advanced search features.

Search with summaries

content_search_spec = discoveryengine.SearchRequest.ContentSearchSpec(
    snippet_spec=discoveryengine.SearchRequest.ContentSearchSpec.SnippetSpec(return_snippet=True),
    summary_spec=discoveryengine.SearchRequest.ContentSearchSpec.SummarySpec(
        summary_result_count=5,
        include_citations=True,
        ignore_adversarial_query=True,
        ignore_non_summary_seeking_query=True,
    ),
)

Requests snippets and citation-backed summaries while filtering adversarial and non-summary queries.

Models & APIs used

  • APIs / services: Vertex AI Search, Cloud Storage, Discovery Engine
  • SDKs / libraries: google-cloud-discoveryengine, google.api_core.client_options, google.cloud.discoveryengine

When to use this

Use this pattern to build a Vertex AI Search app over documents stored in Cloud Storage and query it with summaries and snippets.

Gotchas & caveats

  • Install google-cloud-discoveryengine and restart the notebook runtime before importing the package.
  • Colab requires google.colab authentication; Vertex AI Workbench assumes existing authentication.
  • Application Default Credentials are set with gcloud auth application-default login for the selected project.
  • The datastore name can only contain lowercase letters, numbers, and hyphens.
  • The notebook uses LOCATION = “global”; non-global locations require a regional discoveryengine API endpoint.
  • Datastore and engine creation are long-running operations; datastore creation may exceed the 90 second timeout.
  • New engines may return 404 until they are ready, so wait a few minutes and retry.
  • Enterprise tier is required for extractive answers, and the LLM add-on is required to summarize search results.

Best practices

  • Use ClientOptions with a regional discoveryengine endpoint when location is not global.
  • Use the default_collection parent path for datastore and engine creation.
  • Import documents into the default_branch from Cloud Storage using data_schema=“content”.
  • Use INCREMENTAL reconciliation mode for document import.
  • Enable SEARCH_TIER_ENTERPRISE and SEARCH_ADD_ON_LLM when summaries or advanced LLM features are needed.
  • Include citations in search summaries and ignore adversarial or non-summary-seeking queries.
  • Enable automatic query expansion and spell correction for search requests.