feat: GenAI - GAPIC - Added support for Grounding

feat: Add Retrieval
feat: Add GoogleSearchRetrieval
feat: Add VertexAiSearch
feat: Add Tool.retrieval
feat: Add Tool.google_search_retrieval
feat: Add Candidate.grounding_metadata
PiperOrigin-RevId: 606432042

Author: Ark-kun

google/cloud/aiplatform_v1beta1/__init__.py

@@ -92,9 +92,12 @@
 from .types.content import Content
 from .types.content import FileData
 from .types.content import GenerationConfig
+from .types.content import GroundingAttribution
+from .types.content import GroundingMetadata
 from .types.content import Part
 from .types.content import SafetyRating
 from .types.content import SafetySetting
+from .types.content import Segment
 from .types.content import VideoMetadata
 from .types.content import HarmCategory
 from .types.context import Context
@@ -695,7 +698,10 @@
 from .types.tool import FunctionCall
 from .types.tool import FunctionDeclaration
 from .types.tool import FunctionResponse
+from .types.tool import GoogleSearchRetrieval
+from .types.tool import Retrieval
 from .types.tool import Tool
+from .types.tool import VertexAISearch
 from .types.training_pipeline import FilterSplit
 from .types.training_pipeline import FractionSplit
 from .types.training_pipeline import InputDataConfig
@@ -1064,6 +1070,9 @@
     "GetTensorboardTimeSeriesRequest",
     "GetTrainingPipelineRequest",
     "GetTrialRequest",
+    "GoogleSearchRetrieval",
+    "GroundingAttribution",
+    "GroundingMetadata",
     "HarmCategory",
     "HyperparameterTuningJob",
     "IdMatcher",
@@ -1294,6 +1303,7 @@
     "RestoreDatasetVersionRequest",
     "ResumeModelDeploymentMonitoringJobRequest",
     "ResumeScheduleRequest",
+    "Retrieval",
     "SafetyRating",
     "SafetySetting",
     "SampleConfig",
@@ -1315,6 +1325,7 @@
     "SearchModelDeploymentMonitoringStatsAnomaliesResponse",
     "SearchNearestEntitiesRequest",
     "SearchNearestEntitiesResponse",
+    "Segment",
     "ServiceAccountSpec",
     "SmoothGradConfig",
     "SpecialistPool",
@@ -1410,6 +1421,7 @@
     "UpsertDatapointsResponse",
     "UserActionReference",
     "Value",
+    "VertexAISearch",
     "VideoMetadata",
     "VizierServiceClient",
     "WorkerPoolSpec",

google/cloud/aiplatform_v1beta1/types/__init__.py

@@ -36,9 +36,12 @@
     Content,
     FileData,
     GenerationConfig,
+    GroundingAttribution,
+    GroundingMetadata,
     Part,
     SafetyRating,
     SafetySetting,
+    Segment,
     VideoMetadata,
     HarmCategory,
 )
@@ -784,7 +787,10 @@
     FunctionCall,
     FunctionDeclaration,
     FunctionResponse,
+    GoogleSearchRetrieval,
+    Retrieval,
     Tool,
+    VertexAISearch,
 )
 from .training_pipeline import (
     FilterSplit,
@@ -850,9 +856,12 @@
     "Content",
     "FileData",
     "GenerationConfig",
+    "GroundingAttribution",
+    "GroundingMetadata",
     "Part",
     "SafetyRating",
     "SafetySetting",
+    "Segment",
     "VideoMetadata",
     "HarmCategory",
     "Context",
@@ -1437,7 +1446,10 @@
     "FunctionCall",
     "FunctionDeclaration",
     "FunctionResponse",
+    "GoogleSearchRetrieval",
+    "Retrieval",
     "Tool",
+    "VertexAISearch",
     "FilterSplit",
     "FractionSplit",
     "InputDataConfig",

google/cloud/aiplatform_v1beta1/types/content.py

@@ -39,6 +39,9 @@
         "CitationMetadata",
         "Citation",
         "Candidate",
+        "Segment",
+        "GroundingAttribution",
+        "GroundingMetadata",
     },
 )
 
@@ -503,6 +506,9 @@ class Candidate(proto.Message):
         citation_metadata (google.cloud.aiplatform_v1beta1.types.CitationMetadata):
             Output only. Source attribution of the
             generated content.
+        grounding_metadata (google.cloud.aiplatform_v1beta1.types.GroundingMetadata):
+            Output only. Metadata specifies sources used
+            to ground generated content.
     """
 
     class FinishReason(proto.Enum):
@@ -566,6 +572,125 @@ class FinishReason(proto.Enum):
         number=6,
         message="CitationMetadata",
     )
+    grounding_metadata: "GroundingMetadata" = proto.Field(
+        proto.MESSAGE,
+        number=7,
+        message="GroundingMetadata",
+    )
+
+
+class Segment(proto.Message):
+    r"""Segment of the content.
+
+    Attributes:
+        part_index (int):
+            Output only. The index of a Part object
+            within its parent Content object.
+        start_index (int):
+            Output only. Start index in the given Part,
+            measured in bytes. Offset from the start of the
+            Part, inclusive, starting at zero.
+        end_index (int):
+            Output only. End index in the given Part,
+            measured in bytes. Offset from the start of the
+            Part, exclusive, starting at zero.
+    """
+
+    part_index: int = proto.Field(
+        proto.INT32,
+        number=1,
+    )
+    start_index: int = proto.Field(
+        proto.INT32,
+        number=2,
+    )
+    end_index: int = proto.Field(
+        proto.INT32,
+        number=3,
+    )
+
+
+class GroundingAttribution(proto.Message):
+    r"""Grounding attribution.
+
+    .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields
+
+    Attributes:
+        web (google.cloud.aiplatform_v1beta1.types.GroundingAttribution.Web):
+            Optional. Attribution from the web.
+
+            This field is a member of `oneof`_ ``reference``.
+        segment (google.cloud.aiplatform_v1beta1.types.Segment):
+            Output only. Segment of the content this
+            attribution belongs to.
+        confidence_score (float):
+            Optional. Output only. Confidence score of
+            the attribution. Ranges from 0 to 1. 1 is the
+            most confident.
+
+            This field is a member of `oneof`_ ``_confidence_score``.
+    """
+
+    class Web(proto.Message):
+        r"""Attribution from the web.
+
+        Attributes:
+            uri (str):
+                Output only. URI reference of the
+                attribution.
+            title (str):
+                Output only. Title of the attribution.
+        """
+
+        uri: str = proto.Field(
+            proto.STRING,
+            number=1,
+        )
+        title: str = proto.Field(
+            proto.STRING,
+            number=2,
+        )
+
+    web: Web = proto.Field(
+        proto.MESSAGE,
+        number=3,
+        oneof="reference",
+        message=Web,
+    )
+    segment: "Segment" = proto.Field(
+        proto.MESSAGE,
+        number=1,
+        message="Segment",
+    )
+    confidence_score: float = proto.Field(
+        proto.FLOAT,
+        number=2,
+        optional=True,
+    )
+
+
+class GroundingMetadata(proto.Message):
+    r"""Metadata returned to client when grounding is enabled.
+
+    Attributes:
+        web_search_queries (MutableSequence[str]):
+            Optional. Web search queries for the
+            following-up web search.
+        grounding_attributions (MutableSequence[google.cloud.aiplatform_v1beta1.types.GroundingAttribution]):
+            Optional. List of grounding attributions.
+    """
+
+    web_search_queries: MutableSequence[str] = proto.RepeatedField(
+        proto.STRING,
+        number=1,
+    )
+    grounding_attributions: MutableSequence[
+        "GroundingAttribution"
+    ] = proto.RepeatedField(
+        proto.MESSAGE,
+        number=2,
+        message="GroundingAttribution",
+    )
 
 
 __all__ = tuple(sorted(__protobuf__.manifest))

google/cloud/aiplatform_v1beta1/types/prediction_service.py

@@ -784,8 +784,7 @@ class GenerateContentRequest(proto.Message):
 
             A ``Tool`` is a piece of code that enables the system to
             interact with external systems to perform an action, or set
-            of actions, outside of knowledge and scope of the model. The
-            only supported tool is currently ``Function``
+            of actions, outside of knowledge and scope of the model.
         safety_settings (MutableSequence[google.cloud.aiplatform_v1beta1.types.SafetySetting]):
             Optional. Per request settings for blocking
             unsafe content. Enforced on

google/cloud/aiplatform_v1beta1/types/tool.py

@@ -30,6 +30,9 @@
         "FunctionDeclaration",
         "FunctionCall",
         "FunctionResponse",
+        "Retrieval",
+        "VertexAISearch",
+        "GoogleSearchRetrieval",
     },
 )
 
@@ -39,7 +42,8 @@ class Tool(proto.Message):
 
     A ``Tool`` is a piece of code that enables the system to interact
     with external systems to perform an action, or set of actions,
-    outside of knowledge and scope of the model.
+    outside of knowledge and scope of the model. A Tool object should
+    contain exactly one type of Tool.
 
     Attributes:
         function_declarations (MutableSequence[google.cloud.aiplatform_v1beta1.types.FunctionDeclaration]):
@@ -52,13 +56,32 @@ class Tool(proto.Message):
             function call in the next turn. Based on the function
             responses, Model will generate the final response back to
             the user. Maximum 64 function declarations can be provided.
+        retrieval (google.cloud.aiplatform_v1beta1.types.Retrieval):
+            Optional. System will always execute the
+            provided retrieval tool(s) to get external
+            knowledge to answer the prompt. Retrieval
+            results are presented to the model for
+            generation.
+        google_search_retrieval (google.cloud.aiplatform_v1beta1.types.GoogleSearchRetrieval):
+            Optional. Specialized retrieval tool that is
+            powered by Google search.
     """
 
     function_declarations: MutableSequence["FunctionDeclaration"] = proto.RepeatedField(
         proto.MESSAGE,
         number=1,
         message="FunctionDeclaration",
     )
+    retrieval: "Retrieval" = proto.Field(
+        proto.MESSAGE,
+        number=2,
+        message="Retrieval",
+    )
+    google_search_retrieval: "GoogleSearchRetrieval" = proto.Field(
+        proto.MESSAGE,
+        number=3,
+        message="GoogleSearchRetrieval",
+    )
 
 
 class FunctionDeclaration(proto.Message):
@@ -169,4 +192,71 @@ class FunctionResponse(proto.Message):
     )
 
 
+class Retrieval(proto.Message):
+    r"""Defines a retrieval tool that model can call to access
+    external knowledge.
+
+
+    .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields
+
+    Attributes:
+        vertex_ai_search (google.cloud.aiplatform_v1beta1.types.VertexAISearch):
+            Set to use data source powered by Vertex AI
+            Search.
+
+            This field is a member of `oneof`_ ``source``.
+        disable_attribution (bool):
+            Optional. Disable using the result from this
+            tool in detecting grounding attribution. This
+            does not affect how the result is given to the
+            model for generation.
+    """
+
+    vertex_ai_search: "VertexAISearch" = proto.Field(
+        proto.MESSAGE,
+        number=2,
+        oneof="source",
+        message="VertexAISearch",
+    )
+    disable_attribution: bool = proto.Field(
+        proto.BOOL,
+        number=3,
+    )
+
+
+class VertexAISearch(proto.Message):
+    r"""Retrieve from Vertex AI Search datastore for grounding.
+    See https://cloud.google.com/vertex-ai-search-and-conversation
+
+    Attributes:
+        datastore (str):
+            Required. Fully-qualified Vertex AI Search's
+            datastore resource ID.
+            projects/<>/locations/<>/collections/<>/dataStores/<>
+    """
+
+    datastore: str = proto.Field(
+        proto.STRING,
+        number=1,
+    )
+
+
+class GoogleSearchRetrieval(proto.Message):
+    r"""Tool to retrieve public web data for grounding, powered by
+    Google.
+
+    Attributes:
+        disable_attribution (bool):
+            Optional. Disable using the result from this
+            tool in detecting grounding attribution. This
+            does not affect how the result is given to the
+            model for generation.
+    """
+
+    disable_attribution: bool = proto.Field(
+        proto.BOOL,
+        number=1,
+    )
+
+
 __all__ = tuple(sorted(__protobuf__.manifest))