Update openai_models.py

promptslab · Aug 3, 2023 · c3c9a29 · c3c9a29
1 parent 6c09b6c
commit c3c9a29
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/promptify/models/text2text/api/openai_models.py b/promptify/models/text2text/api/openai_models.py
@@ -184,9 +184,15 @@ def _default_max_tokens(self, model_name: str) -> int:
     def model_output_raw(self, response: Dict) -> Dict:
         data = {}
         if self.model_type == "chat_models":
-            data["text"] = response["choices"][0]["message"]["content"].strip(" \n")
+            try:
+              data["text"] = response["choices"][0]["message"]["content"].strip(" \n")
+            except Exception as e:
+              data["text"] = response[0]["choices"][0]["message"]["content"].strip(" \n")
         elif self.model_type == "completion_models":
+          try:
             data["text"] = response["choices"][0]["text"]
+          except Exception as e:
+            data["text"] = response[0]["choices"][0]["text"]
 
         data["usage"] = dict(response["usage"])
         return data