Merge pull request #204 from IINemo/fix/openai_chat

OpenAI chat try-except added, prompt fixed

Author: rvashurin

src/lm_polygraph/utils/openai_chat.py

@@ -1,10 +1,15 @@
 import openai
 import json
 import os
+import time
+import logging
 
 from filelock import FileLock
 
 
+log = logging.getLogger()
+
+
 class OpenAIChat:
     """
     Allows for the implementation of a singleton class to chat with OpenAI model for dataset marking.
@@ -50,12 +55,11 @@ def ask(self, message: str) -> str:
                     "Please specify OPENAI_KEY in environment parameters."
                 )
             messages = [
-                {"role": "system", "content": "You are a intelligent assistant."},
+                {"role": "system", "content": "You are an intelligent assistant."},
                 {"role": "user", "content": message},
             ]
-            chat = openai.ChatCompletion.create(
-                model=self.openai_model, messages=messages
-            )
+            chat = self._send_request(messages)
+
             reply = chat.choices[0].message.content
 
             # add reply to cache
@@ -76,3 +80,19 @@ def ask(self, message: str) -> str:
             return ""
 
         return reply
+
+    def _send_request(self, messages):
+        sleep_time_values = (5, 10, 30, 60, 120)
+        for i in range(len(sleep_time_values)):
+            try:
+                return openai.ChatCompletion.create(
+                    model=self.openai_model, messages=messages
+                )
+            except Exception as e:
+                sleep_time = sleep_time_values[i]
+                log.info(
+                    f"Request to OpenAI failed with exception: {e}. Retry #{i}/5 after {sleep_time} seconds."
+                )
+                time.sleep(sleep_time)
+
+        return openai.ChatCompletion.create(model=self.openai_model, messages=messages)