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

Ark-kun · copybara-github · commit 459ba86396ab · 2023-06-29T18:01:51.000-07:00
PiperOrigin-RevId: 544512632
diff --git a/tests/unit/aiplatform/test_language_models.py b/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",
diff --git a/vertexai/language_models/_language_models.py b/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")

Original file line number	Diff line number	Diff line change
`@@ -799,7 +799,8 @@ def send_message(`
`799`	`799`	`)`
`800`	`800`
`801`	`801`	`prediction = prediction_response.predictions[0]`
`802`		`- safety_attributes = prediction["safetyAttributes"]`
	`802`	`+ # ! Note: For chat models, the safetyAttributes is a list.`
	`803`	`+ safety_attributes = prediction["safetyAttributes"][0]`
`803`	`804`	`response_obj = TextGenerationResponse(`
`804`	`805`	`text=prediction["candidates"][0]["content"]`
`805`	`806`	`if prediction.get("candidates")`