feat: add endpoint field to BatchPredictionJob for bring-your-own-endpoint batch prediction

google/cloud/aiplatform_v1/types/tool.py

@@ -102,6 +102,11 @@ class Tool(proto.Message):
             user queries. The search results will be
             grounded on Parallel.ai and presented to the
             model for response generation
+        exa_ai_search (google.cloud.aiplatform_v1.types.Tool.ExaAiSearch):
+            Optional. Uses Exa.ai to search for
+            information to answer user queries. The search
+            results will be grounded on Exa.ai and presented
+            to the model for response generation
         code_execution (google.cloud.aiplatform_v1.types.Tool.CodeExecution):
             Optional. CodeExecution tool type.
             Enables the model to execute code as part of
@@ -222,6 +227,28 @@ class ParallelAiSearch(proto.Message):
             message=struct_pb2.Struct,
         )
 
+    class ExaAiSearch(proto.Message):
+        r"""ExaAiSearch tool type.
+        A tool that uses the Exa.ai search engine for grounding.
+
+        Attributes:
+            api_key (str):
+                Required. The API key for ExaAiSearch.
+            custom_configs (google.protobuf.struct_pb2.Struct):
+                Optional. This field can be used to pass any
+                parameter from the Exa.ai Search API.
+        """
+
+        api_key: str = proto.Field(
+            proto.STRING,
+            number=1,
+        )
+        custom_configs: struct_pb2.Struct = proto.Field(
+            proto.MESSAGE,
+            number=2,
+            message=struct_pb2.Struct,
+        )
+
     class CodeExecution(proto.Message):
         r"""Tool that executes code generated by the model, and automatically
         returns the result to the model.
@@ -308,6 +335,11 @@ class Environment(proto.Enum):
         number=13,
         message=ParallelAiSearch,
     )
+    exa_ai_search: ExaAiSearch = proto.Field(
+        proto.MESSAGE,
+        number=15,
+        message=ExaAiSearch,
+    )
     code_execution: CodeExecution = proto.Field(
         proto.MESSAGE,
         number=4,

google/cloud/aiplatform_v1beta1/types/batch_prediction_job.py

@@ -90,6 +90,12 @@ class BatchPredictionJob(proto.Message):
             prediction without requiring uploading to model registry.
             Exactly one of model and unmanaged_container_model must be
             set.
+        endpoint (str):
+            For Bring-Your-Own-Endpoint (BYOE), the name of the Endpoint
+            resource that produces the predictions via this job, must
+            share the same ancestor Location. Exactly one of model,
+            unmanaged_container_model, or endpoint must be set. Example:
+            ``projects/193595526740/locations/us-central1/endpoints/4203439000301600768``
         input_config (google.cloud.aiplatform_v1beta1.types.BatchPredictionJob.InputConfig):
             Required. Input configuration of the instances on which
             predictions are performed. The schema of any single instance
@@ -637,6 +643,10 @@ class OutputInfo(proto.Message):
             message=gca_unmanaged_container_model.UnmanagedContainerModel,
         )
     )
+    endpoint: str = proto.Field(
+        proto.STRING,
+        number=41,
+    )
     input_config: InputConfig = proto.Field(
         proto.MESSAGE,
         number=4,

google/cloud/aiplatform_v1beta1/types/tool.py

@@ -103,6 +103,11 @@ class Tool(proto.Message):
             user queries. The search results will be
             grounded on Parallel.ai and presented to the
             model for response generation
+        exa_ai_search (google.cloud.aiplatform_v1beta1.types.Tool.ExaAiSearch):
+            Optional. Uses Exa.ai to search for
+            information to answer user queries. The search
+            results will be grounded on Exa.ai and presented
+            to the model for response generation
         code_execution (google.cloud.aiplatform_v1beta1.types.Tool.CodeExecution):
             Optional. CodeExecution tool type.
             Enables the model to execute code as part of
@@ -223,6 +228,28 @@ class ParallelAiSearch(proto.Message):
             message=struct_pb2.Struct,
         )
 
+    class ExaAiSearch(proto.Message):
+        r"""ExaAiSearch tool type.
+        A tool that uses the Exa.ai search engine for grounding.
+
+        Attributes:
+            api_key (str):
+                Required. The API key for ExaAiSearch.
+            custom_configs (google.protobuf.struct_pb2.Struct):
+                Optional. This field can be used to pass any
+                parameter from the Exa.ai Search API.
+        """
+
+        api_key: str = proto.Field(
+            proto.STRING,
+            number=1,
+        )
+        custom_configs: struct_pb2.Struct = proto.Field(
+            proto.MESSAGE,
+            number=2,
+            message=struct_pb2.Struct,
+        )
+
     class CodeExecution(proto.Message):
         r"""Tool that executes code generated by the model, and automatically
         returns the result to the model.
@@ -309,6 +336,11 @@ class Environment(proto.Enum):
         number=13,
         message=ParallelAiSearch,
     )
+    exa_ai_search: ExaAiSearch = proto.Field(
+        proto.MESSAGE,
+        number=15,
+        message=ExaAiSearch,
+    )
     code_execution: CodeExecution = proto.Field(
         proto.MESSAGE,
         number=4,

samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json

@@ -8,7 +8,7 @@
     ],
     "language": "PYTHON",
     "name": "google-cloud-aiplatform",
-    "version": "1.163.0"
+    "version": "0.0.0"
   },
   "snippets": [
     {

samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json

@@ -8,7 +8,7 @@
     ],
     "language": "PYTHON",
     "name": "google-cloud-aiplatform",
-    "version": "1.163.0"
+    "version": "0.0.0"
   },
   "snippets": [
     {