vertex_ai_search_text_query_with_discovery_engine.py

1from typing import Any
2
3from google.api_core.client_options import ClientOptions
4from google.cloud import discoveryengine_v1 as discoveryengine
5
6# TODO(developer): Update these variables before running the sample.
7project_id = "YOUR_PROJECT_ID"
8location = "global"          # Values: "global", "us", "eu"
9engine_id = "YOUR_ENGINE_ID" # The ID of the search engine/app
10search_query = "What is text search?"
11
12def search_sample(
13    project_id: str,
14    location: str,
15    engine_id: str,
16    search_query: str,
17) -> Any:
18    #  API endpoint for the Discovery Engine
19    client_options = (
20        ClientOptions(api_endpoint=f"{location}-discoveryengine.googleapis.com")
21        if location != "global"
22        else None
23    )
24
25    # Create a client
26    client = discoveryengine.SearchServiceClient(client_options=client_options)
27
28    # The full resource name of the search engine serving config
29    # e.g. projects/{project_id}/locations/{location}/collections/default_collection/engines/{engine_id}/servingConfigs/default_serving_config
30    serving_config = client.serving_config_path(
31        project=project_id,
32        location=location,
33        data_store=engine_id,
34        serving_config="default_serving_config",
35    )
36
37    request = discoveryengine.SearchRequest(
38        serving_config=serving_config,
39        query=search_query,
40        page_size=10,
41    )
42
43    response = client.search(request)
44
45    print("Search results:")
46    for result in response.results:
47        print(result.document.derived_struct_data)
48
49    return response
50
51if __name__ == "__main__":
52    search_sample(project_id, location, engine_id, search_query)