refactor: (GenAI) Reorganized Grounding Samples (Group B)

generative_ai/gemini_grounding_example.py

@@ -11,6 +11,7 @@
 # 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.
+# TODO: Delete this file after approval /grounding/web_example.py & /grounding/vais_example.py
 import os
 
 from vertexai.generative_models import GenerationResponse

generative_ai/grounding.py

@@ -11,6 +11,7 @@
 # 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.
+# TODO: Delete this file after approval /grounding/palm_example.py
 import os
 
 from typing import Optional

generative_ai/grounding/palm_example.py

@@ -0,0 +1,72 @@
+# 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
+#
+#    https://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 typing import Optional
+
+from vertexai.language_models import TextGenerationResponse
+
+PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+
+
+def grounding(
+    data_store_location: Optional[str] = None,
+    data_store_id: Optional[str] = None,
+) -> TextGenerationResponse:
+    """Grounding example with a Large Language Model"""
+    # [START generativeaionvertexai_grounding]
+    import vertexai
+
+    from vertexai.language_models import GroundingSource, TextGenerationModel
+
+    # TODO(developer): Update and un-comment below line
+    # PROJECT_ID = "your-project-id"
+    vertexai.init(project=PROJECT_ID, location="us-central1")
+
+    # TODO developer - override these parameters as needed:
+    parameters = {
+        "temperature": 0.1,  # Temperature controls the degree of randomness in token selection.
+        "max_output_tokens": 256,  # Token limit determines the maximum amount of text output.
+        "top_p": 0.8,  # Tokens are selected from most probable to least until the sum of their probabilities equals the top_p value.
+        "top_k": 20,  # A top_k of 1 means the selected token is the most probable among all tokens.
+    }
+
+    model = TextGenerationModel.from_pretrained("text-bison@002")
+
+    # TODO(developer): Update and un-comment below lines
+    # data_store_id = "datastore_123456789012345"
+    # data_store_location = "global"
+    if data_store_id and data_store_location:
+        # Use Vertex AI Search data store
+        grounding_source = GroundingSource.VertexAISearch(
+            data_store_id=data_store_id, location=data_store_location
+        )
+    else:
+        # Use Google Search for grounding (Private Preview)
+        grounding_source = GroundingSource.WebSearch()
+
+    response = model.predict(
+        "What are the price, available colors, and storage size options of a Pixel Tablet?",
+        grounding_source=grounding_source,
+        **parameters,
+    )
+    print(f"Response from Model: {response.text}")
+    print(f"Grounding Metadata: {response.grounding_metadata}")
+    # [END generativeaionvertexai_grounding]
+
+    return response
+
+
+if __name__ == "__main__":
+    grounding(data_store_id="data-store_1234567890123", data_store_location="global")

generative_ai/grounding/test_grounding.py

@@ -0,0 +1,50 @@
+# 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
+#
+#    https://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 backoff
+
+from google.api_core.exceptions import ResourceExhausted
+
+import palm_example
+import pytest
+
+import vais_example
+import web_example
+
+
+@backoff.on_exception(backoff.expo, ResourceExhausted, max_time=10)
+def test_grounding() -> None:
+    data_store_id = "test-search-engine_1689960780551"
+    response = palm_example.grounding(
+        data_store_location="global",
+        data_store_id=data_store_id,
+    )
+    assert response
+    assert response.text
+    assert response.grounding_metadata
+
+
+def test_gemini_grounding_vais_example() -> None:
+    response = vais_example.generate_text_with_grounding_vertex_ai_search(
+        "grounding-test-datastore"
+    )
+    assert response
+
+
+@pytest.mark.skip(
+    "Unable to test Google Search grounding due to allowlist restrictions."
+)
+def test_gemini_grounding_web_example() -> None:
+    response = web_example.generate_text_with_grounding_web()
+    assert response

generative_ai/grounding/vais_example.py

@@ -0,0 +1,68 @@
+# 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
+#
+#    https://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 vertexai.generative_models import GenerationResponse
+
+PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+
+
+def generate_text_with_grounding_vertex_ai_search(
+    data_store_id: str,
+) -> GenerationResponse:
+    # [START generativeaionvertexai_gemini_grounding_with_vais]
+    import vertexai
+
+    from vertexai.preview.generative_models import (
+        GenerationConfig,
+        GenerativeModel,
+        Tool,
+        grounding,
+    )
+
+    # TODO(developer): Update and un-comment below lines
+    # PROJECT_ID = "your-project-id"
+    # data_store_id = "your-data-store-id"
+
+    vertexai.init(project=PROJECT_ID, location="us-central1")
+
+    model = GenerativeModel("gemini-1.5-flash-001")
+
+    tool = Tool.from_retrieval(
+        grounding.Retrieval(
+            grounding.VertexAISearch(
+                datastore=data_store_id,
+                project=PROJECT_ID,
+                location="global",
+            )
+        )
+    )
+
+    prompt = "How do I make an appointment to renew my driver's license?"
+    response = model.generate_content(
+        prompt,
+        tools=[tool],
+        generation_config=GenerationConfig(
+            temperature=0.0,
+        ),
+    )
+
+    print(response.text)
+
+    # [END generativeaionvertexai_gemini_grounding_with_vais]
+    return response
+
+
+if __name__ == "__main__":
+    generate_text_with_grounding_vertex_ai_search("data-store_1234567890123")

generative_ai/grounding/web_example.py

@@ -0,0 +1,59 @@
+# 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
+#
+#    https://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 vertexai.generative_models import GenerationResponse
+
+PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+
+
+def generate_text_with_grounding_web() -> GenerationResponse:
+    # [START generativeaionvertexai_gemini_grounding_with_web]
+    import vertexai
+
+    from vertexai.generative_models import (
+        GenerationConfig,
+        GenerativeModel,
+        Tool,
+        grounding,
+    )
+
+    # TODO(developer): Update and un-comment below line
+    # PROJECT_ID = "your-project-id"
+    vertexai.init(project=PROJECT_ID, location="us-central1")
+
+    model = GenerativeModel("gemini-1.5-flash-001")
+
+    # Use Google Search for grounding
+    tool = Tool.from_google_search_retrieval(grounding.GoogleSearchRetrieval())
+
+    prompt = "When is the next total solar eclipse in US?"
+    response = model.generate_content(
+        prompt,
+        tools=[tool],
+        generation_config=GenerationConfig(
+            temperature=0.0,
+        ),
+    )
+
+    print(response.text)
+    # Example response:
+    # The next total solar eclipse visible from the contiguous United States will be on **August 23, 2044**.
+
+    # [END generativeaionvertexai_gemini_grounding_with_web]
+    return response
+
+
+if __name__ == "__main__":
+    generate_text_with_grounding_web()

generative_ai/test_gemini_examples.py

@@ -127,6 +127,7 @@ def test_gemini_chat_example() -> None:
     assert any([_ in text for _ in ("hi", "hello", "greeting")])
 
 
+# TODO: Delete this file after approval /grounding/web_example.py
 @pytest.mark.skip(
     "Unable to test Google Search grounding due to allowlist restrictions."
 )
@@ -135,6 +136,7 @@ def test_gemini_grounding_web_example() -> None:
     assert response
 
 
+# TODO: Delete this file after approval /grounding/vais_example.py
 def test_gemini_grounding_vais_example() -> None:
     response = gemini_grounding_example.generate_text_with_grounding_vertex_ai_search(
         "grounding-test-datastore"