samples(documentai, discoveryengine): Add Samples for Vertex AI Search Standalone RAG APIs (GoogleCloudPlatform#11613)

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: holtskinner

discoveryengine/requirements.txt

@@ -1,2 +1,2 @@
-google-cloud-discoveryengine==0.11.10
+google-cloud-discoveryengine==0.11.11
 google-api-core==2.17.1

discoveryengine/standalone_apis_sample.py

@@ -0,0 +1,125 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from google.cloud import discoveryengine_v1alpha as discoveryengine
+
+
+def check_grounding_sample(
+    project_id: str,
+) -> discoveryengine.CheckGroundingResponse:
+    # [START genappbuilder_check_grounding]
+    from google.cloud import discoveryengine_v1alpha as discoveryengine
+
+    # TODO(developer): Uncomment these variables before running the sample.
+    # project_id = "YOUR_PROJECT_ID"
+
+    client = discoveryengine.GroundedGenerationServiceClient()
+
+    # The full resource name of the grounding config.
+    # Format: projects/{project_id}/locations/{location}/groundingConfigs/default_grounding_config
+    grounding_config = client.grounding_config_path(
+        project=project_id,
+        location="global",
+        grounding_config="default_grounding_config",
+    )
+
+    request = discoveryengine.CheckGroundingRequest(
+        grounding_config=grounding_config,
+        answer_candidate="Titanic was directed by James Cameron. It was released in 1997.",
+        facts=[
+            discoveryengine.GroundingFact(
+                fact_text=(
+                    "Titanic is a 1997 American epic romantic disaster movie. It was directed, written,"
+                    " and co-produced by James Cameron. The movie is about the 1912 sinking of the"
+                    " RMS Titanic. It stars Kate Winslet and Leonardo DiCaprio. The movie was released"
+                    " on December 19, 1997. It received positive critical reviews. The movie won 11 Academy"
+                    " Awards, and was nominated for fourteen total Academy Awards."
+                ),
+                attributes={"author": "Simple Wikipedia"},
+            ),
+            discoveryengine.GroundingFact(
+                fact_text=(
+                    'James Cameron\'s "Titanic" is an epic, action-packed romance'
+                    "set against the ill-fated maiden voyage of the R.M.S. Titanic;"
+                    "the pride and joy of the White Star Line and, at the time,"
+                    "the largest moving object ever built. "
+                    'She was the most luxurious liner of her era -- the "ship of dreams" -- '
+                    "which ultimately carried over 1,500 people to their death in the "
+                    "ice cold waters of the North Atlantic in the early hours of April 15, 1912."
+                ),
+                attributes={"author": "Simple Wikipedia"},
+            ),
+        ],
+        grounding_spec=discoveryengine.CheckGroundingSpec(citation_threshold=0.6),
+    )
+
+    response = client.check_grounding(request=request)
+
+    # Handle the response
+    print(response)
+    # [END genappbuilder_check_grounding]
+
+    return response
+
+
+def rank_sample(
+    project_id: str,
+) -> discoveryengine.RankResponse:
+    # [START genappbuilder_rank]
+    from google.cloud import discoveryengine_v1alpha as discoveryengine
+
+    # TODO(developer): Uncomment these variables before running the sample.
+    # project_id = "YOUR_PROJECT_ID"
+
+    client = discoveryengine.RankServiceClient()
+
+    # The full resource name of the ranking config.
+    # Format: projects/{project_id}/locations/{location}/rankingConfigs/default_ranking_config
+    ranking_config = client.ranking_config_path(
+        project=project_id,
+        location="global",
+        ranking_config="default_ranking_config",
+    )
+    request = discoveryengine.RankRequest(
+        ranking_config=ranking_config,
+        model="semantic-ranker-512@latest",
+        top_n=10,
+        query="What is Google Gemini?",
+        records=[
+            discoveryengine.RankingRecord(
+                id="1",
+                title="Gemini",
+                content="The Gemini zodiac symbol often depicts two figures standing side-by-side.",
+            ),
+            discoveryengine.RankingRecord(
+                id="2",
+                title="Gemini",
+                content="Gemini is a cutting edge large language model created by Google.",
+            ),
+            discoveryengine.RankingRecord(
+                id="3",
+                title="Gemini Constellation",
+                content="Gemini is a constellation that can be seen in the night sky.",
+            ),
+        ],
+    )
+
+    response = client.rank(request=request)
+
+    # Handle the response
+    print(response)
+    # [END genappbuilder_rank]
+
+    return response

discoveryengine/standalone_apis_sample_test.py

@@ -0,0 +1,34 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+
+from discoveryengine import standalone_apis_sample
+
+project_id = os.environ["GOOGLE_CLOUD_PROJECT"]
+
+
+def test_check_grounding():
+    response = standalone_apis_sample.check_grounding_sample(project_id)
+    assert response
+    assert response.support_score
+    assert response.cited_chunks
+    assert response.claims
+
+
+def test_rank():
+    response = standalone_apis_sample.rank_sample(project_id)
+    assert response
+    assert response.records

documentai/snippets/handle_response_sample_test.py

@@ -219,3 +219,18 @@ def test_process_document_summarizer(capsys):
     ]
     for expected_string in expected_strings:
         assert expected_string in out
+
+
+def test_process_document_layout():
+    document = handle_response_sample_v1beta3.process_document_layout_sample(
+        project_id=os.environ["GOOGLE_CLOUD_PROJECT"],
+        location="us",
+        processor_id="85b02a52f356f564",
+        processor_version="pretrained",
+        file_path="resources/superconductivity.pdf",
+        mime_type="application/pdf",
+    )
+
+    assert document
+    assert document.document_layout
+    assert document.chunked_document

documentai/snippets/handle_response_sample_v1beta3.py

@@ -30,8 +30,51 @@
 # mime_type = "application/pdf" # Refer to https://cloud.google.com/document-ai/docs/file-types for supported file types
 
 # [END documentai_process_custom_extractor_document]
+# [END documentai_process_summarizer_document]
+
 
+# [START documentai_process_layout_document]
+def process_document_layout_sample(
+    project_id: str,
+    location: str,
+    processor_id: str,
+    processor_version: str,
+    file_path: str,
+    mime_type: str,
+) -> documentai.Document:
+    process_options = documentai.ProcessOptions(
+        layout_config=documentai.ProcessOptions.LayoutConfig(
+            chunking_config=documentai.ProcessOptions.LayoutConfig.ChunkingConfig(
+                chunk_size=1000,
+                include_ancestor_headings=True,
+                breakpoint_percentile_threshold=90,
+            )
+        )
+    )
 
+    document = process_document(
+        project_id,
+        location,
+        processor_id,
+        processor_version,
+        file_path,
+        mime_type,
+        process_options=process_options,
+    )
+
+    print("Document Layout Blocks")
+    for block in document.document_layout.blocks:
+        print(block)
+
+    print("Document Chunks")
+    for chunk in document.chunked_document.chunks:
+        print(chunk)
+
+    # [END documentai_process_layout_document]
+    return document
+
+
+# [START documentai_process_summarizer_document]
 def process_document_summarizer_sample(
     project_id: str,
     location: str,
@@ -175,6 +218,7 @@ def print_entity(entity: documentai.Document.Entity) -> None:
         print(f"    * Normalized Value: {repr(normalized_value)}")
 
 
+# [START documentai_process_layout_document]
 def process_document(
     project_id: str,
     location: str,
@@ -217,5 +261,6 @@ def process_document(
     return result.document
 
 
+# [END documentai_process_layout_document]
 # [END documentai_process_summarizer_document]
 # [END documentai_process_custom_extractor_document]

documentai/snippets/requirements.txt

@@ -1,2 +1,2 @@
-google-cloud-documentai==2.20.0
-google-cloud-storage==2.9.0
+google-cloud-documentai==2.27.0
+google-cloud-storage==2.16.0