Upload check_openai.py

utils_groupclassification/check_openai.py

@@ -4,6 +4,7 @@ from logging import getLogger
 
 from openai import OpenAI
 from openai import AzureOpenAI
+from openai import OpenAIError
 import fitz
 
 import requests
@@ -55,26 +56,41 @@ def generate_check_(reference):
     client = OpenAI(
         api_key=api_key,
     )
-    response = client.chat.completions.create(
-        model="gpt-3.5-turbo",
-        messages=[
-            {
-                "role": "system",
-                "content": system_prompt,
-            },
-            {
-                "role": "user",
-                "content": reference,
-            },
-        ],
-        functions=[{"name": "generate_queries", "parameters": json_schema}],
-        function_call={"name": "generate_queries"},
-        temperature=0.0,
-        top_p=0.0,
-    )
-    output = response.choices[0].message.function_call.arguments
-    time.sleep(1)
-    return output
+    retries = 0
+    max_retries = 100
+    delay = 5
+
+    while retries < max_retries:
+        try:
+            response = client.chat.completions.create(
+                model="gpt-3.5-turbo",
+                messages=[
+                    {
+                        "role": "system",
+                        "content": system_prompt,
+                    },
+                    {
+                        "role": "user",
+                        "content": reference,
+                    },
+                ],
+                functions=[{"name": "generate_queries", "parameters": json_schema}],
+                function_call={"name": "generate_queries"},
+                temperature=0.0,
+                top_p=0.0,
+            )
+            output = response.choices[0].message.function_call.arguments
+            time.sleep(1)
+            return output
+        except OpenAIError as e:
+            print(f"Error occurred: {e}. Retrying in {delay} seconds...")
+            retries += 1
+            time.sleep(delay)
+        except Exception as e:
+            print(f"Unexpected error: {e}. Retrying in {delay} seconds...")
+            retries += 1
+            time.sleep(delay)
+    raise RuntimeError("Maximum retries exceeded. Could not get a valid response.")
 
 
 def ch(company_name, reference):