Extend answer check + retry mechanism to oai impl

pskl · Nov 17, 2023 · 255fda0 · 255fda0
1 parent 75dc17b
commit 255fda0
Showing 1 changed file with 18 additions and 8 deletions.
diff --git a/lib/implementations/openai_impl.py b/lib/implementations/openai_impl.py
@@ -1,13 +1,23 @@
 import openai
 import os
+import time
 
 class OpenaiImpl:
   def ask_question(self, question, prompt, model):
-      response = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY")).chat.completions.create(
-          model=model,
-          messages=[
-              {"role": "system", "content": prompt},
-              {"role": "user", "content": question}
-          ]
-      )
-      return response.choices[0].message.content
+    for i in range(15):
+        try:
+            response = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY")).chat.completions.create(
+                model=model,
+                messages=[
+                    {"role": "system", "content": prompt},
+                    {"role": "user", "content": question}
+                ]
+            )
+            int_response = int(response.choices[0].message.content)
+            return int_response
+        except ValueError:
+            time.sleep(1)
+            print("Retrying the same question\n")
+            continue
+    else:
+        raise ValueError("POST request is not returning integer response. Model is misbehaving, needs further inspection.")