samples(discoveryengine): Add Grounded Generation API Samples

discoveryengine/requirements-test.txt

@@ -1,2 +1,3 @@
 pytest==8.2.0
-google-api-core==2.17.1
+google-api-core==2.21.0
+google-cloud-resource-manager==1.12.5

discoveryengine/requirements.txt

@@ -1,2 +1,2 @@
-google-cloud-discoveryengine==0.12.1
-google-api-core==2.19.0
+google-cloud-discoveryengine==0.12.3
+google-api-core==2.21.0

discoveryengine/site_search_engine_sample.py

@@ -29,7 +29,8 @@ def create_target_site(
     # project_id = "YOUR_PROJECT_ID"
     # location = "YOUR_LOCATION" # Values: "global"
     # data_store_id = "YOUR_DATA_STORE_ID"
-    # uri_pattern = "https://cloud.google.com/generative-ai-app-builder/docs/*"
+    # NOTE: Do not include http or https protocol in the URI pattern
+    # uri_pattern = "cloud.google.com/generative-ai-app-builder/docs/*"
 
     #  For more information, refer to:
     # https://cloud.google.com/generative-ai-app-builder/docs/locations#specify_a_multi-region_for_your_data_store

discoveryengine/site_search_engine_sample_test.py

@@ -28,7 +28,7 @@ def test_create_target_site():
         project_id,
         location,
         data_store_id,
-        uri_pattern="https://cloud.google.com/generative-ai-app-builder/docs/*",
+        uri_pattern="cloud.google.com/generative-ai-app-builder/docs/*",
     )
     assert response, response
     match = re.search(r"\/targetSites\/([^\/]+)", response.name)

discoveryengine/standalone_apis_sample.py

@@ -13,14 +13,14 @@
 # limitations under the License.
 #
 
-from google.cloud import discoveryengine_v1alpha as discoveryengine
+from google.cloud import discoveryengine_v1 as discoveryengine
 
 
 def check_grounding_sample(
     project_id: str,
 ) -> discoveryengine.CheckGroundingResponse:
     # [START genappbuilder_check_grounding]
-    from google.cloud import discoveryengine_v1alpha as discoveryengine
+    from google.cloud import discoveryengine_v1 as discoveryengine
 
     # TODO(developer): Uncomment these variables before running the sample.
     # project_id = "YOUR_PROJECT_ID"
@@ -78,7 +78,7 @@ def rank_sample(
     project_id: str,
 ) -> discoveryengine.RankResponse:
     # [START genappbuilder_rank]
-    from google.cloud import discoveryengine_v1alpha as discoveryengine
+    from google.cloud import discoveryengine_v1 as discoveryengine
 
     # TODO(developer): Uncomment these variables before running the sample.
     # project_id = "YOUR_PROJECT_ID"
@@ -123,3 +123,183 @@ def rank_sample(
     # [END genappbuilder_rank]
 
     return response
+
+
+def grounded_generation_inline_vais_sample(
+    project_number: str,
+    engine_id: str,
+) -> discoveryengine.GenerateGroundedContentResponse:
+    # [START genappbuilder_grounded_generation_inline_vais]
+    from google.cloud import discoveryengine_v1 as discoveryengine
+
+    # TODO(developer): Uncomment these variables before running the sample.
+    # project_number = "YOUR_PROJECT_NUMBER"
+    # engine_id = "YOUR_ENGINE_ID"
+
+    client = discoveryengine.GroundedGenerationServiceClient()
+
+    request = discoveryengine.GenerateGroundedContentRequest(
+        # The full resource name of the location.
+        # Format: projects/{project_number}/locations/{location}
+        location=client.common_location_path(project=project_number, location="global"),
+        generation_spec=discoveryengine.GenerateGroundedContentRequest.GenerationSpec(
+            model_id="gemini-1.5-flash",
+        ),
+        # Conversation between user and model
+        contents=[
+            discoveryengine.GroundedGenerationContent(
+                role="user",
+                parts=[
+                    discoveryengine.GroundedGenerationContent.Part(
+                        text="How did Google do in 2020? Where can I find BigQuery docs?"
+                    )
+                ],
+            )
+        ],
+        system_instruction=discoveryengine.GroundedGenerationContent(
+            parts=[
+                discoveryengine.GroundedGenerationContent.Part(
+                    text="Add a smiley emoji after the answer."
+                )
+            ],
+        ),
+        # What to ground on.
+        grounding_spec=discoveryengine.GenerateGroundedContentRequest.GroundingSpec(
+            grounding_sources=[
+                discoveryengine.GenerateGroundedContentRequest.GroundingSource(
+                    inline_source=discoveryengine.GenerateGroundedContentRequest.GroundingSource.InlineSource(
+                        grounding_facts=[
+                            discoveryengine.GroundingFact(
+                                fact_text=(
+                                    "The BigQuery documentation can be found at https://cloud.google.com/bigquery/docs/introduction"
+                                ),
+                                attributes={
+                                    "title": "BigQuery Overview",
+                                    "uri": "https://cloud.google.com/bigquery/docs/introduction",
+                                },
+                            ),
+                        ]
+                    ),
+                ),
+                discoveryengine.GenerateGroundedContentRequest.GroundingSource(
+                    search_source=discoveryengine.GenerateGroundedContentRequest.GroundingSource.SearchSource(
+                        # The full resource name of the serving config for a Vertex AI Search App
+                        serving_config=f"projects/{project_number}/locations/global/collections/default_collection/engines/{engine_id}/servingConfigs/default_search",
+                    ),
+                ),
+            ]
+        ),
+    )
+    response = client.generate_grounded_content(request)
+
+    # Handle the response
+    print(response)
+    # [END genappbuilder_grounded_generation_inline_vais]
+
+    return response
+
+
+def grounded_generation_google_search_sample(
+    project_number: str,
+) -> discoveryengine.GenerateGroundedContentResponse:
+    # [START genappbuilder_grounded_generation_google_search]
+    from google.cloud import discoveryengine_v1 as discoveryengine
+
+    # TODO(developer): Uncomment these variables before running the sample.
+    # project_number = "YOUR_PROJECT_NUMBER"
+
+    client = discoveryengine.GroundedGenerationServiceClient()
+
+    request = discoveryengine.GenerateGroundedContentRequest(
+        # The full resource name of the location.
+        # Format: projects/{project_number}/locations/{location}
+        location=client.common_location_path(project=project_number, location="global"),
+        generation_spec=discoveryengine.GenerateGroundedContentRequest.GenerationSpec(
+            model_id="gemini-1.5-flash",
+        ),
+        # Conversation between user and model
+        contents=[
+            discoveryengine.GroundedGenerationContent(
+                role="user",
+                parts=[
+                    discoveryengine.GroundedGenerationContent.Part(
+                        text="How much is Google stock?"
+                    )
+                ],
+            )
+        ],
+        system_instruction=discoveryengine.GroundedGenerationContent(
+            parts=[
+                discoveryengine.GroundedGenerationContent.Part(text="Be comprehensive.")
+            ],
+        ),
+        # What to ground on.
+        grounding_spec=discoveryengine.GenerateGroundedContentRequest.GroundingSpec(
+            grounding_sources=[
+                discoveryengine.GenerateGroundedContentRequest.GroundingSource(
+                    google_search_source=discoveryengine.GenerateGroundedContentRequest.GroundingSource.GoogleSearchSource(
+                        # Optional: For Dynamic Retrieval
+                        dynamic_retrieval_config=discoveryengine.GenerateGroundedContentRequest.DynamicRetrievalConfiguration(
+                            predictor=discoveryengine.GenerateGroundedContentRequest.DynamicRetrievalConfiguration.DynamicRetrievalPredictor(
+                                threshold=0.7
+                            )
+                        )
+                    )
+                ),
+            ]
+        ),
+    )
+    response = client.generate_grounded_content(request)
+
+    # Handle the response
+    print(response)
+    # [END genappbuilder_grounded_generation_google_search]
+
+    return response
+
+
+def grounded_generation_streaming_sample(
+    project_number: str,
+) -> discoveryengine.GenerateGroundedContentResponse:
+    # [START genappbuilder_grounded_generation_streaming]
+    from google.cloud import discoveryengine_v1 as discoveryengine
+
+    # TODO(developer): Uncomment these variables before running the sample.
+    # project_id = "YOUR_PROJECT_ID"
+
+    client = discoveryengine.GroundedGenerationServiceClient()
+
+    request = discoveryengine.GenerateGroundedContentRequest(
+        # The full resource name of the location.
+        # Format: projects/{project_number}/locations/{location}
+        location=client.common_location_path(project=project_number, location="global"),
+        generation_spec=discoveryengine.GenerateGroundedContentRequest.GenerationSpec(
+            model_id="gemini-1.5-flash",
+        ),
+        # Conversation between user and model
+        contents=[
+            discoveryengine.GroundedGenerationContent(
+                role="user",
+                parts=[
+                    discoveryengine.GroundedGenerationContent.Part(
+                        text="Summarize how to delete a data store in Vertex AI Agent Builder?"
+                    )
+                ],
+            )
+        ],
+        grounding_spec=discoveryengine.GenerateGroundedContentRequest.GroundingSpec(
+            grounding_sources=[
+                discoveryengine.GenerateGroundedContentRequest.GroundingSource(
+                    google_search_source=discoveryengine.GenerateGroundedContentRequest.GroundingSource.GoogleSearchSource()
+                ),
+            ]
+        ),
+    )
+    responses = client.stream_generate_grounded_content(iter([request]))
+
+    for response in responses:
+        # Handle the response
+        print(response)
+    # [END genappbuilder_grounded_generation_streaming]
+
+    return response

discoveryengine/standalone_apis_sample_test.py

@@ -17,6 +17,8 @@
 
 from discoveryengine import standalone_apis_sample
 
+from google.cloud import resourcemanager_v3
+
 project_id = os.environ["GOOGLE_CLOUD_PROJECT"]
 
 
@@ -32,3 +34,27 @@ def test_rank():
     response = standalone_apis_sample.rank_sample(project_id)
     assert response
     assert response.records
+
+
+def test_grounded_generation_inline_vais_sample():
+    # Grounded Generation requires Project Number
+    client = resourcemanager_v3.ProjectsClient()
+    project = client.get_project(name=client.project_path(project_id))
+    project_number = client.parse_project_path(project.name)["project"]
+
+    response = standalone_apis_sample.grounded_generation_inline_vais_sample(
+        project_number, engine_id="test-search-engine_1689960780551"
+    )
+    assert response
+
+
+def test_grounded_generation_google_search_sample():
+    # Grounded Generation requires Project Number
+    client = resourcemanager_v3.ProjectsClient()
+    project = client.get_project(name=client.project_path(project_id))
+    project_number = client.parse_project_path(project.name)["project"]
+
+    response = standalone_apis_sample.grounded_generation_google_search_sample(
+        project_number
+    )
+    assert response