fix: PrivateEndpoint predict method should set parameters in request payload.

vertex-sdk-bot · copybara-github · commit c7526a1d44b5 · 2025-02-20T10:23:10.000-08:00
PiperOrigin-RevId: 729160042
diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py
@@ -3614,7 +3614,7 @@ def predict(
         Example usage:
             PSA based private endpoint:
 
-            response = my_private_endpoint.predict(instances=[...])
+            response = my_private_endpoint.predict(instances=[...], parameters={...})
             my_predictions = response.predictions
 
             PSC based private endpoint:
@@ -3680,7 +3680,7 @@ def predict(
             response = self._http_request(
                 method="POST",
                 url=self.predict_http_uri,
-                body=json.dumps({"instances": instances}),
+                body=json.dumps({"instances": instances, "parameters": parameters}),
                 headers={"Content-Type": "application/json"},
             )
             prediction_response = json.loads(response.data)
@@ -3717,7 +3717,7 @@ def predict(
             response = self._http_request(
                 method="POST",
                 url=url,
-                body=json.dumps({"instances": instances}),
+                body=json.dumps({"instances": instances, "parameters": parameters}),
                 headers=headers,
             )
 
diff --git a/tests/unit/aiplatform/test_endpoints.py b/tests/unit/aiplatform/test_endpoints.py
@@ -3525,7 +3525,7 @@ def test_psa_predict(self, predict_private_endpoint_mock):
         predict_private_endpoint_mock.assert_called_once_with(
             method="POST",
             url="",
-            body='{"instances": [[1.0, 2.0, 3.0], [1.0, 3.0, 4.0]]}',
+            body='{"instances": [[1.0, 2.0, 3.0], [1.0, 3.0, 4.0]], "parameters": {"param": 3.0}}',
             headers={"Content-Type": "application/json"},
         )
 
@@ -3552,7 +3552,7 @@ def test_psc_predict(self, predict_private_endpoint_mock):
         predict_private_endpoint_mock.assert_called_once_with(
             method="POST",
             url=f"https://{_TEST_ENDPOINT_OVERRIDE}/v1/projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/endpoints/{_TEST_ID}:predict",
-            body='{"instances": [[1.0, 2.0, 3.0], [1.0, 3.0, 4.0]]}',
+            body='{"instances": [[1.0, 2.0, 3.0], [1.0, 3.0, 4.0]], "parameters": {"param": 3.0}}',
             headers={
                 "Content-Type": "application/json",
                 "Authorization": "Bearer None",
