fix: Handle case when no metadata is returned from model.predict

PiperOrigin-RevId: 607414090
googleapis · Feb 15, 2024 · 59e2bca · 59e2bca
1 parent cfd96d8
commit 59e2bca
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py
@@ -1564,7 +1564,7 @@ def predict(
             json_response = raw_predict_response.json()
             return Prediction(
                 predictions=json_response["predictions"],
-                metadata=json_response["metadata"],
+                metadata=json_response.get("metadata"),
                 deployed_model_id=raw_predict_response.headers[
                     _RAW_PREDICT_DEPLOYED_MODEL_ID_KEY
                 ],
@@ -1582,7 +1582,10 @@ def predict(
                 parameters=parameters,
                 timeout=timeout,
             )
-            metadata = json_format.MessageToDict(prediction_response._pb.metadata)
+            if prediction_response._pb.metadata:
+                metadata = json_format.MessageToDict(prediction_response._pb.metadata)
+            else:
+                metadata = None
 
             return Prediction(
                 predictions=[
@@ -1643,7 +1646,10 @@ async def predict_async(
             parameters=parameters,
             timeout=timeout,
         )
-        metadata = json_format.MessageToDict(prediction_response._pb.metadata)
+        if prediction_response._pb.metadata:
+            metadata = json_format.MessageToDict(prediction_response._pb.metadata)
+        else:
+            metadata = None
 
         return Prediction(
             predictions=[