FIxed stuff

src/utils/openaicaller.py

@@ -37,13 +37,14 @@ class bcolors:
     BOLD = '\033[1m'
     UNDERLINE = '\033[4m'
 
-chat_models = ["gpt-4", "gpt-4-32k", "gpt-3.5-turbo", "gpt-3.5-turbo-16k"]
+chat_models = ["gpt-4", "gpt-4-32k", "gpt-3.5-turbo", "gpt-3.5-turbo-16k", "gpt-3.5-turbo-0613"]
 text_models = ["text-davinci-003", "text-davinci-002", "text-curie-001", "text-babbage-001", "text-ada-001"]
 
 models_max_tokens = {
     "gpt-4": 8_192,
     "gpt-4-32k": 32_768,
     "gpt-3.5-turbo": 4_096,
+    "gpt-3.5-turbo-0613": 4_096,
     "gpt-3.5-turbo-16k": 16_384,
     "text-davinci-003": 4_097,
     "text-davinci-002": 4_097,
@@ -54,17 +55,17 @@ models_max_tokens = {
 
 class openai_caller:
     def __init__(self, api_key=None) -> None:
-        pass
-    def set_api_key(self, key):
-        openai_module.api_key = key
-    async def generate_response(self, **kwargs):
-        if kwargs['model'] in chat_models:
-            return await self.chat_generate(**kwargs)
-        elif kwargs['model'] in text_models:
+        self.api_key = api_key
+    async def generate_response(self, error_call=None, **kwargs):
+        if error_call is None:
+            error_call = lambda x: 2 # do nothing
+        if kwargs.get("model", "") in chat_models:
+            return await self.chat_generate(error_call, **kwargs)
+        elif kwargs.get("engine", "") in text_models:
             raise NotImplementedError("Text models are not supported yet")
         else:
             raise ValueError("Model not found")
-    async def chat_generate(self, **kwargs):
+    async def chat_generate(self, recall_func, **kwargs):
         tokens = await num_tokens_from_messages(kwargs['messages'], kwargs['model'])
         model_max_tokens = models_max_tokens[kwargs['model']]
         while tokens > model_max_tokens:
@@ -73,36 +74,68 @@ class openai_caller:
             tokens = await num_tokens_from_messages(kwargs['messages'], kwargs['model'])
         i = 0
         response = None
+        kwargs['api_key'] = self.api_key
         while i < 10:
             try:
-                response = await openai_module.ChatCompletion.acreate(**kwargs)
+                response = await openai_module.ChatCompletion.acreate(
+                    **kwargs
+                    )
                 break
-            except APIError:
+            except APIError as e:
+                print(f"\n\n{bcolors.BOLD}{bcolors.WARNING}APIError. This is not your fault. Retrying...{bcolors.ENDC}")
+                await recall_func("`An APIError occurred. This is not your fault. Retrying...`")
                 await asyncio.sleep(10)
+                await recall_func()
                 i += 1
-            except Timeout:
+            except Timeout as e:
+                print(f"\n\n{bcolors.BOLD}{bcolors.WARNING}The request timed out. Retrying...{bcolors.ENDC}")
+                await recall_func("`The request timed out. Retrying...`")
                 await asyncio.sleep(10)
+                await recall_func()
                 i += 1
-            except RateLimitError:
+            except RateLimitError as e:
+                print(f"\n\n{bcolors.BOLD}{bcolors.WARNING}RateLimitError. You are being rate limited. Retrying...{bcolors.ENDC}")
+                await recall_func("`You are being rate limited. Retrying...`")
                 await asyncio.sleep(10)
+                await recall_func()
                 i += 1
             except APIConnectionError as e:
-                print(e)
                 print(f"\n\n{bcolors.BOLD}{bcolors.FAIL}APIConnectionError. There is an issue with your internet connection. Please check your connection.{bcolors.ENDC}")
+                await recall_func()
                 raise e
             except InvalidRequestError as e:
-                print(e)
                 print(f"\n\n{bcolors.BOLD}{bcolors.FAIL}InvalidRequestError. Please check your request.{bcolors.ENDC}")
+                await recall_func()
                 raise e
             except AuthenticationError as e:
-                print(e)
-                print(f"\n\n{bcolors.BOLD}{bcolors.FAIL}AuthenticationError. Please check your API key.{bcolors.ENDC}")
+                print(f"\n\n{bcolors.BOLD}{bcolors.FAIL}AuthenticationError. Please check your API key and if needed, also your organization ID.{bcolors.ENDC}")
+                await recall_func("`AuthenticationError. Please check your API key.`")
                 raise e
-            except ServiceUnavailableError:
+            except ServiceUnavailableError as e:
+                print(f"\n\n{bcolors.BOLD}{bcolors.WARNING}ServiceUnavailableError. The OpenAI API is not responding. Retrying...{bcolors.ENDC}")
+                await recall_func("`The OpenAI API is not responding. Retrying...`")
                 await asyncio.sleep(10)
+                await recall_func()
                 i += 1
             finally:
                 if i == 10:
                     print(f"\n\n{bcolors.BOLD}{bcolors.FAIL}OpenAI API is not responding. Please try again later.{bcolors.ENDC}")
                     raise TimeoutError("OpenAI API is not responding. Please try again later.")
-        return response # type: ignore
+        return response # type: ignore
+    
+##testing
+if __name__ == "__main__":
+    async def main():
+        openai = openai_caller(api_key="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")
+        response = await openai.generate_response(
+            model="gpt-3.5-turbo",
+            messages=[{"role":"user", "content":"ping"}],
+            max_tokens=5,
+            temperature=0.7,
+            top_p=1,
+            frequency_penalty=0,
+            presence_penalty=0,
+            stop=["\n", " Human:", " AI:"]
+        )
+        print(response)
+    asyncio.run(main())