From 645dfa8bc28518b9db48e8c06f524d5a1443533a Mon Sep 17 00:00:00 2001
From: Azalea <22280294+hykilpikonna@users.noreply.github.com>
Date: Mon, 9 Mar 2026 00:43:24 -0400
Subject: [PATCH] [F] Fix dict keys

---
 utils_ai.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/utils_ai.py b/utils_ai.py
index dfe5e81..7e9eb47 100644
--- a/utils_ai.py
+++ b/utils_ai.py
@@ -38,10 +38,16 @@ def _call_openai_with_prompt(prompt_id: str, prompt_version: str, input_text: st
     try:
         # If response is a Pydantic model
         if hasattr(response, 'output') and response.output:
-            return response.output[0].content[0].text
+            for out_msg in response.output:
+                if getattr(out_msg, 'type', '') == 'message':
+                    return out_msg.content[0].text
+            return response.output[-1].content[0].text
         # If response is a dictionary
         elif isinstance(response, dict) and "output" in response:
-            return response["output"][0]["content"][0]["text"]
+            for out_msg in response["output"]:
+                if out_msg.get("type") == "message":
+                    return out_msg["content"][0]["text"]
+            return response["output"][-1]["content"][0]["text"]
         # Fallback for choices (if API changes slightly)
         elif hasattr(response, 'choices') and response.choices:
             return response.choices[0].message.content