Skip to content

Commit

Permalink
feat: GenAI - Grounding - Released Google Web Search retriever to GA
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 632335942
  • Loading branch information
Ark-kun authored and Copybara-Service committed May 10, 2024
1 parent 12c147b commit 32b030a
Show file tree
Hide file tree
Showing 4 changed files with 43 additions and 2 deletions.
19 changes: 18 additions & 1 deletion tests/unit/vertexai/test_generative_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -868,7 +868,7 @@ def test_conversion_methods(self, generative_models: generative_models):
attribute="generate_content",
new=mock_generate_content,
)
def test_generate_content_grounding_google_search_retriever(self):
def test_generate_content_grounding_google_search_retriever_preview(self):
model = preview_generative_models.GenerativeModel("gemini-pro")
google_search_retriever_tool = (
preview_generative_models.Tool.from_google_search_retrieval(
Expand All @@ -882,6 +882,23 @@ def test_generate_content_grounding_google_search_retriever(self):
)
assert response.text

@mock.patch.object(
target=prediction_service.PredictionServiceClient,
attribute="generate_content",
new=mock_generate_content,
)
def test_generate_content_grounding_google_search_retriever(self):
model = generative_models.GenerativeModel("gemini-pro")
google_search_retriever_tool = (
generative_models.Tool.from_google_search_retrieval(
generative_models.grounding.GoogleSearchRetrieval()
)
)
response = model.generate_content(
"Why is sky blue?", tools=[google_search_retriever_tool]
)
assert response.text

@mock.patch.object(
target=prediction_service.PredictionServiceClient,
attribute="generate_content",
Expand Down
2 changes: 2 additions & 0 deletions vertexai/generative_models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
ResponseValidationError,
SafetySetting,
Tool,
grounding,
)

__all__ = [
Expand All @@ -50,4 +51,5 @@
"ResponseValidationError",
"SafetySetting",
"Tool",
"grounding",
]
22 changes: 22 additions & 0 deletions vertexai/generative_models/_generative_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -2012,6 +2012,28 @@ def __repr__(self):
class grounding: # pylint: disable=invalid-name
"""Grounding namespace."""

__module__ = "vertexai.generative_models"

def __init__(self):
raise RuntimeError("This class must not be instantiated.")

class GoogleSearchRetrieval:
r"""Tool to retrieve public web data for grounding, powered by
Google Search.
"""

def __init__(self):
"""Initializes a Google Search Retrieval tool.
"""
self._raw_google_search_retrieval = gapic_tool_types.GoogleSearchRetrieval()


class preview_grounding: # pylint: disable=invalid-name
"""Grounding namespace (preview)."""

__name__ = "grounding"
__module__ = "vertexai.preview.generative_models"

def __init__(self):
raise RuntimeError("This class must not be instantiated.")

Expand Down
2 changes: 1 addition & 1 deletion vertexai/preview/generative_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
# We just want to re-export certain classes
# pylint: disable=g-multiple-import,g-importing-member
from vertexai.generative_models._generative_models import (
grounding,
preview_grounding as grounding,
_PreviewGenerativeModel,
_PreviewChatSession,
GenerationConfig,
Expand Down

0 comments on commit 32b030a

Please sign in to comment.