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

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,113 @@
+# 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 typing import List
+
+from google.cloud import discoveryengine_v1alpha as discoveryengine
+
+
+def check_grounding_sample(
+    project_id: str,
+) -> List[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,
+) -> List[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 2023 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,25 @@ def test_process_document_summarizer(capsys):
     ]
     for expected_string in expected_strings:
         assert expected_string in out
+
+
+def test_process_document_layout():
+    location = "us"
+    project_id = os.environ["GOOGLE_CLOUD_PROJECT"]
+    processor_id = "85b02a52f356f564"
+    processor_version = "pretrained"
+    file_path = "resources/superconductivity.pdf"
+    mime_type = "application/pdf"
+
+    document = handle_response_sample_v1beta3.process_document_layout_sample(
+        project_id=project_id,
+        location=location,
+        processor_id=processor_id,
+        processor_version=processor_version,
+        file_path=file_path,
+        mime_type=mime_type,
+    )
+
+    assert document
+    assert document.document_layout
+    assert document.chunked_document

documentai/snippets/handle_response_sample_v1beta3.py

@@ -30,8 +30,52 @@
 # 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,
+            )
+        )
+    )
+
+    # Online processing request to Document AI
+    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 +219,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 +262,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