chat: handle not connected better

rmackay9 · Dec 8, 2023 · 20ffacc · 20ffacc
1 parent 5c01a63
commit 20ffacc
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 10 deletions.
diff --git a/MAVProxy/modules/mavproxy_chat/chat_openai.py b/MAVProxy/modules/mavproxy_chat/chat_openai.py
@@ -20,14 +20,38 @@ def __init__(self, mpstate):
         # keep reference to mpstate
         self.mpstate = mpstate
 
+        # initialise OpenAI connection
+        self.client = None
+        self.assistant = None
+
+        # default assistant id
+        self.assistant_id = "asst_XqQN1m1JXtIROLqkjcuNYzZ0"
+
+    # check connection to OpenAI assistant and connect if necessary
+    # returns True if connection is good, False if not
+    def check_connection(self):
         # create connection object
-        self.client = OpenAI()
+        if self.client is None:
+            try:
+                self.client = OpenAI()
+            except:
+                print("chat: failed to connect to OpenAI")
+                return False
+
+        # check connection again just to be sure
+        if self.client is None:
+            print("chat: failed to connect to OpenAI")
+            return False
 
         # get assistant id
-        self.assistant_id = "asst_XqQN1m1JXtIROLqkjcuNYzZ0"
-        self.assistant = self.client.beta.assistants.retrieve(self.assistant_id)
         if self.assistant is None:
-            print("chat: failed to find OpenAI assistant")
+            self.assistant = self.client.beta.assistants.retrieve(self.assistant_id)
+            if self.assistant is None:
+                print("chat: failed to connect to OpenAI assistant")
+                return False
+
+        # if we got this far the connection must be good
+        return True
 
     # set the OpenAI API key
     def set_api_key(self, api_key_str):
@@ -38,9 +62,9 @@ def send_to_assistant(self, text):
         # debug
         print("Sending text to assistant: " + text)
 
-        # check connection to assistant
-        if self.assistant is None:
-            return "chat: failed to connect to OpenAI assistant"
+        # check connection
+        if not self.check_connection():
+            return "chat: failed to connect to OpenAI"
 
         # create new thread
         self.assistant_thread = self.client.beta.threads.create(

diff --git a/MAVProxy/modules/mavproxy_chat/chat_voice_to_text.py b/MAVProxy/modules/mavproxy_chat/chat_voice_to_text.py
@@ -9,9 +9,19 @@
 from openai import OpenAI
 
 class chat_voice_to_text():
-    def __init__(self):
+    # check connection to OpenAI assistant and connect if necessary
+    # returns True if connection is good, False if not
+    def check_connection(self):
         # create connection object
-        self.client = OpenAI()
+        if self.client is None:
+            try:
+                self.client = OpenAI()
+            except:
+                print("chat: failed to connect to OpenAI")
+                return False
+
+        # return True if connected
+        return self.client is not None
 
     # record audio from microphone
     # returns filename of recording or None if failed
@@ -57,9 +67,12 @@ def record_audio(self):
     # convert audio to text
     # returns transcribed text on success or None if failed
     def convert_audio_to_text(self, audio_filename):
+        # check connection
+        if not self.check_connection():
+            return None
+
         # Process with Whisper
         audio_file = open(audio_filename, "rb")
-        client = OpenAI()
         transcript = client.audio.transcriptions.create(
             model="whisper-1", 
             file=audio_file,