feat: Add Grounding feature to PredictionService.GenerateContent

feat: Add Retrieval
feat: Add GoogleSearchRetrieval
feat: Add VertexAiSearch
feat: Add Tool.retrieval
feat: Add Tool.google_search_retrieval
feat: Add Candidate.grounding_metadata

docs: A comment for message `Tool` is changed
docs: A comment for field `tools` in message `.google.cloud.aiplatform.v1beta1.GenerateContentRequest` is changed
docs: A comment for field `tools` in message `.google.cloud.aiplatform.v1beta1.GenerateContentRequest` is changed
PiperOrigin-RevId: 605743031

google/cloud/aiplatform/v1beta1/content.proto

@@ -291,4 +291,60 @@ message Candidate {
   // Output only. Source attribution of the generated content.
   CitationMetadata citation_metadata = 6
       [(google.api.field_behavior) = OUTPUT_ONLY];
+
+  // Output only. Metadata specifies sources used to ground generated content.
+  GroundingMetadata grounding_metadata = 7
+      [(google.api.field_behavior) = OUTPUT_ONLY];
+}
+
+// Segment of the content.
+message Segment {
+  // Output only. The index of a Part object within its parent Content object.
+  int32 part_index = 1 [(google.api.field_behavior) = OUTPUT_ONLY];
+
+  // Output only. Start index in the given Part, measured in bytes. Offset from
+  // the start of the Part, inclusive, starting at zero.
+  int32 start_index = 2 [(google.api.field_behavior) = OUTPUT_ONLY];
+
+  // Output only. End index in the given Part, measured in bytes. Offset from
+  // the start of the Part, exclusive, starting at zero.
+  int32 end_index = 3 [(google.api.field_behavior) = OUTPUT_ONLY];
+}
+
+// Grounding attribution.
+message GroundingAttribution {
+  // Attribution from the web.
+  message Web {
+    // Output only. URI reference of the attribution.
+    string uri = 1 [(google.api.field_behavior) = OUTPUT_ONLY];
+
+    // Output only. Title of the attribution.
+    string title = 2 [(google.api.field_behavior) = OUTPUT_ONLY];
+  }
+
+  oneof reference {
+    // Optional. Attribution from the web.
+    Web web = 3 [(google.api.field_behavior) = OPTIONAL];
+  }
+
+  // Output only. Segment of the content this attribution belongs to.
+  Segment segment = 1 [(google.api.field_behavior) = OUTPUT_ONLY];
+
+  // Optional. Output only. Confidence score of the attribution. Ranges from 0
+  // to 1. 1 is the most confident.
+  optional float confidence_score = 2 [
+    (google.api.field_behavior) = OPTIONAL,
+    (google.api.field_behavior) = OUTPUT_ONLY
+  ];
+}
+
+// Metadata returned to client when grounding is enabled.
+message GroundingMetadata {
+  // Optional. Web search queries for the following-up web search.
+  repeated string web_search_queries = 1
+      [(google.api.field_behavior) = OPTIONAL];
+
+  // Optional. List of grounding attributions.
+  repeated GroundingAttribution grounding_attributions = 2
+      [(google.api.field_behavior) = OPTIONAL];
 }

google/cloud/aiplatform/v1beta1/tool.proto

@@ -32,7 +32,8 @@ option ruby_package = "Google::Cloud::AIPlatform::V1beta1";
 //
 // 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.
+// knowledge and scope of the model. A Tool object should contain exactly
+// one type of Tool.
 message Tool {
   // Optional. One or more function declarations to be passed to the model along
   // with the current user query. Model may decide to call a subset of these
@@ -44,6 +45,15 @@ message Tool {
   // provided.
   repeated FunctionDeclaration function_declarations = 1
       [(google.api.field_behavior) = OPTIONAL];
+
+  // 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.
+  Retrieval retrieval = 2 [(google.api.field_behavior) = OPTIONAL];
+
+  // Optional. Specialized retrieval tool that is powered by Google search.
+  GoogleSearchRetrieval google_search_retrieval = 3
+      [(google.api.field_behavior) = OPTIONAL];
 }
 
 // Structured representation of a function declaration as defined by the
@@ -102,3 +112,32 @@ message FunctionResponse {
   // Required. The function response in JSON object format.
   google.protobuf.Struct response = 2 [(google.api.field_behavior) = REQUIRED];
 }
+
+// Defines a retrieval tool that model can call to access external knowledge.
+message Retrieval {
+  oneof source {
+    // Set to use data source powered by Vertex AI Search.
+    VertexAISearch vertex_ai_search = 2;
+  }
+
+  // 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.
+  bool disable_attribution = 3 [(google.api.field_behavior) = OPTIONAL];
+}
+
+// Retrieve from Vertex AI Search datastore for grounding.
+// See https://cloud.google.com/vertex-ai-search-and-conversation
+message VertexAISearch {
+  // Required. Fully-qualified Vertex AI Search's datastore resource ID.
+  // projects/<>/locations/<>/collections/<>/dataStores/<>
+  string datastore = 1 [(google.api.field_behavior) = REQUIRED];
+}
+
+// Tool to retrieve public web data for grounding, powered by Google.
+message GoogleSearchRetrieval {
+  // 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.
+  bool disable_attribution = 1 [(google.api.field_behavior) = OPTIONAL];
+}