diff --git a/mock-responses/streaming-failure-finish-reason-safety.txt b/mock-responses/streaming-failure-finish-reason-safety.txt
index 05e0936..b318cba 100644
--- a/mock-responses/streaming-failure-finish-reason-safety.txt
+++ b/mock-responses/streaming-failure-finish-reason-safety.txt
@@ -1,2 +1,2 @@
-data: {"candidates": [{"content": {"parts": [{"text": "<redacted>"}],"role": "model"},"finishReason": "SAFETY","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "HIGH"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}}
+data: {"candidates": [{"content": {"parts": [{"text": "<redacted>"}],"role": "model"},"finishReason": "SAFETY","safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE","probabilityScore": 0.40429688,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.18554688},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE","probabilityScore": 0.15820313,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.030639648},{"category": "HARM_CATEGORY_HARASSMENT","probability": "LOW","blocked": true,"probabilityScore": 0.55078125,"severity": "HARM_SEVERITY_LOW","severityScore": 0.24804688},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE","probabilityScore": 0.17675781,"severity": "HARM_SEVERITY_NEGLIGIBLE","severityScore": 0.17871094}]}],"usageMetadata": {"promptTokenCount": 10,"candidatesTokenCount": 66,"totalTokenCount": 76},"modelVersion": "gemini-1.5-flash"}
 
diff --git a/mock-responses/unary-failure-finish-reason-safety-no-content.json b/mock-responses/unary-failure-finish-reason-safety-no-content.json
index 03958d4..4d0e1f7 100644
--- a/mock-responses/unary-failure-finish-reason-safety-no-content.json
+++ b/mock-responses/unary-failure-finish-reason-safety-no-content.json
@@ -1,46 +1,43 @@
 {
-    "candidates": [
+  "candidates": [
+    {
+      "finishReason": "SAFETY",
+      "safetyRatings": [
         {
-            "finishReason": "SAFETY",
-            "index": 0,
-            "safetyRatings": [
-                {
-                    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-                    "probability": "NEGLIGIBLE"
-                },
-                {
-                    "category": "HARM_CATEGORY_HATE_SPEECH",
-                    "probability": "HIGH"
-                },
-                {
-                    "category": "HARM_CATEGORY_HARASSMENT",
-                    "probability": "NEGLIGIBLE"
-                },
-                {
-                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-                    "probability": "NEGLIGIBLE"
-                }
-            ]
+          "category": "HARM_CATEGORY_HATE_SPEECH",
+          "probability": "NEGLIGIBLE",
+          "probabilityScore": 0.3984375,
+          "severity": "HARM_SEVERITY_LOW",
+          "severityScore": 0.21582031
+        },
+        {
+          "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+          "probability": "NEGLIGIBLE",
+          "probabilityScore": 0.14941406,
+          "severity": "HARM_SEVERITY_NEGLIGIBLE",
+          "severityScore": 0.02331543
+        },
+        {
+          "category": "HARM_CATEGORY_HARASSMENT",
+          "probability": "LOW",
+          "blocked": true,
+          "probabilityScore": 0.61328125,
+          "severity": "HARM_SEVERITY_LOW",
+          "severityScore": 0.31835938
+        },
+        {
+          "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+          "probability": "NEGLIGIBLE",
+          "probabilityScore": 0.13476563,
+          "severity": "HARM_SEVERITY_NEGLIGIBLE",
+          "severityScore": 0.12109375
         }
-    ],
-    "promptFeedback": {
-        "safetyRatings": [
-            {
-                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-                "probability": "NEGLIGIBLE"
-            },
-            {
-                "category": "HARM_CATEGORY_HATE_SPEECH",
-                "probability": "NEGLIGIBLE"
-            },
-            {
-                "category": "HARM_CATEGORY_HARASSMENT",
-                "probability": "NEGLIGIBLE"
-            },
-            {
-                "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-                "probability": "NEGLIGIBLE"
-            }
-        ]
+      ]
     }
+  ],
+  "usageMetadata": {
+    "promptTokenCount": 8,
+    "totalTokenCount": 8
+  },
+  "modelVersion": "gemini-1.5-flash"
 }
diff --git a/mock-responses/unary-failure-finish-reason-safety.json b/mock-responses/unary-failure-finish-reason-safety.json
index 111e33d..ca27481 100644
--- a/mock-responses/unary-failure-finish-reason-safety.json
+++ b/mock-responses/unary-failure-finish-reason-safety.json
@@ -10,45 +10,42 @@
         "role": "model"
       },
       "finishReason": "SAFETY",
-      "index": 0,
       "safetyRatings": [
         {
-          "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-          "probability": "NEGLIGIBLE"
+          "category": "HARM_CATEGORY_HATE_SPEECH",
+          "probability": "NEGLIGIBLE",
+          "probabilityScore": 0.3984375,
+          "severity": "HARM_SEVERITY_LOW",
+          "severityScore": 0.21582031
         },
         {
-          "category": "HARM_CATEGORY_HATE_SPEECH",
-          "probability": "HIGH"
+          "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+          "probability": "NEGLIGIBLE",
+          "probabilityScore": 0.14941406,
+          "severity": "HARM_SEVERITY_NEGLIGIBLE",
+          "severityScore": 0.02331543
         },
         {
           "category": "HARM_CATEGORY_HARASSMENT",
-          "probability": "NEGLIGIBLE"
+          "probability": "LOW",
+          "blocked": true,
+          "probabilityScore": 0.61328125,
+          "severity": "HARM_SEVERITY_LOW",
+          "severityScore": 0.31835938
         },
         {
-          "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-          "probability": "NEGLIGIBLE"
+          "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+          "probability": "NEGLIGIBLE",
+          "probabilityScore": 0.13476563,
+          "severity": "HARM_SEVERITY_NEGLIGIBLE",
+          "severityScore": 0.12109375
         }
       ]
     }
   ],
-  "promptFeedback": {
-    "safetyRatings": [
-      {
-        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-        "probability": "NEGLIGIBLE"
-      },
-      {
-        "category": "HARM_CATEGORY_HATE_SPEECH",
-        "probability": "NEGLIGIBLE"
-      },
-      {
-        "category": "HARM_CATEGORY_HARASSMENT",
-        "probability": "NEGLIGIBLE"
-      },
-      {
-        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-        "probability": "NEGLIGIBLE"
-      }
-    ]
-  }
+  "usageMetadata": {
+    "promptTokenCount": 8,
+    "totalTokenCount": 8
+  },
+  "modelVersion": "gemini-1.5-flash"
 }