From 6e2cade4a163da1b85a419b3188ca6f1b8ca3591 Mon Sep 17 00:00:00 2001
From: Nicolas van Kempen <nvankemp@gmail.com>
Date: Mon, 8 Jan 2024 17:34:38 +0000
Subject: [PATCH] Fix OpenAI issues

---
 src/chatdbg/chatdbg_utils.py | 39 ++++++++++++++++++++++++------------
 src/chatdbg/chatdbg_why.py   | 33 ++++++++++++++++++++----------
 2 files changed, 48 insertions(+), 24 deletions(-)

diff --git a/src/chatdbg/chatdbg_utils.py b/src/chatdbg/chatdbg_utils.py
index acdcf81..67db13c 100644
--- a/src/chatdbg/chatdbg_utils.py
+++ b/src/chatdbg/chatdbg_utils.py
@@ -75,18 +75,31 @@ def explain(source_code: str, traceback: str, exception: str, really_run=True) -
         return
 
     try:
-        completion = openai.ChatCompletion.create(
-            model=model,
-            request_timeout=30,
-            messages=[{"role": "user", "content": user_prompt}],
-        )
-        text = completion.choices[0].message.content
-        input_tokens = completion.usage.prompt_tokens
-        output_tokens = completion.usage.completion_tokens
-        cost = llm_utils.calculate_cost(input_tokens, output_tokens, model)
-        text += f"\n(Total cost: approximately ${cost:.2f} USD.)"
-        print(llm_utils.word_wrap_except_code_blocks(text))
-    except openai.error.AuthenticationError:
-        print("You need a valid OpenAI key to use ChatDBG.")
+        client = openai.OpenAI(timeout=30)
+    except openai.OpenAIError:
+        print("You need an OpenAI key to use this tool.")
         print("You can get a key here: https://platform.openai.com/api-keys")
         print("Set the environment variable OPENAI_API_KEY to your key value.")
+        return
+
+    try:
+        completion = client.chat.completions.create(
+            model=model, messages=[{"role": "user", "content": user_prompt}]
+        )
+    except openai.NotFoundError:
+        print(f"'{model}' either does not exist or you do not have access to it.")
+        return
+    except openai.RateLimitError:
+        print("You have exceeded a rate limit or have no remaining funds.")
+        return
+    except openai.APITimeoutError:
+        print("The OpenAI API timed out.")
+        return
+
+    text = completion.choices[0].message.content
+    print(llm_utils.word_wrap_except_code_blocks(text))
+
+    input_tokens = completion.usage.prompt_tokens
+    output_tokens = completion.usage.completion_tokens
+    cost = llm_utils.calculate_cost(input_tokens, output_tokens, model)
+    print(f"\n(Total cost: approximately ${cost:.2f} USD.)")
diff --git a/src/chatdbg/chatdbg_why.py b/src/chatdbg/chatdbg_why.py
index c3ac9df..c7be2f9 100644
--- a/src/chatdbg/chatdbg_why.py
+++ b/src/chatdbg/chatdbg_why.py
@@ -68,15 +68,26 @@ def why(self, arg):
         return
 
     try:
-        completion = openai.ChatCompletion.create(
-            model=model,
-            request_timeout=30,
-            messages=[{"role": "user", "content": user_prompt}],
-        )
-        text = completion.choices[0].message.content
-        print(llm_utils.word_wrap_except_code_blocks(text))
-    except openai.error.AuthenticationError:
-        print(
-            "You need a valid OpenAI key to use ChatDBG. You can get a key here: https://openai.com/api/"
-        )
+        client = openai.OpenAI(timeout=30)
+    except openai.OpenAIError:
+        print("You need an OpenAI key to use this tool.")
+        print("You can get a key here: https://platform.openai.com/api-keys")
         print("Set the environment variable OPENAI_API_KEY to your key value.")
+        return
+
+    try:
+        completion = client.chat.completions.create(
+            model=model, messages=[{"role": "user", "content": user_prompt}]
+        )
+    except openai.NotFoundError:
+        print(f"'{model}' either does not exist or you do not have access to it.")
+        return
+    except openai.RateLimitError:
+        print("You have exceeded a rate limit or have no remaining funds.")
+        return
+    except openai.APITimeoutError:
+        print("The OpenAI API timed out.")
+        return
+
+    text = completion.choices[0].message.content
+    print(llm_utils.word_wrap_except_code_blocks(text))