fix: LLM - Fixed the chat models failing due to safetyAttributes format

PiperOrigin-RevId: 544512632

Author: Ark-kun

tests/unit/aiplatform/test_language_models.py

@@ -164,11 +164,13 @@
 }
 
 _TEST_CHAT_GENERATION_PREDICTION1 = {
-    "safetyAttributes": {
-        "scores": [],
-        "blocked": False,
-        "categories": [],
-    },
+    "safetyAttributes": [
+        {
+            "scores": [],
+            "blocked": False,
+            "categories": [],
+        }
+    ],
     "candidates": [
         {
             "author": "1",
@@ -177,11 +179,13 @@
     ],
 }
 _TEST_CHAT_GENERATION_PREDICTION2 = {
-    "safetyAttributes": {
-        "scores": [],
-        "blocked": False,
-        "categories": [],
-    },
+    "safetyAttributes": [
+        {
+            "scores": [],
+            "blocked": False,
+            "categories": [],
+        }
+    ],
     "candidates": [
         {
             "author": "1",

vertexai/language_models/_language_models.py

@@ -799,7 +799,8 @@ def send_message(
         )
 
         prediction = prediction_response.predictions[0]
-        safety_attributes = prediction["safetyAttributes"]
+        # ! Note: For chat models, the safetyAttributes is a list.
+        safety_attributes = prediction["safetyAttributes"][0]
         response_obj = TextGenerationResponse(
             text=prediction["candidates"][0]["content"]
             if prediction.get("candidates")