chat: enabled chat streaming feature

Co-authored-by: Randy Mackay <[email protected]>

Author: 2 people

MAVProxy/modules/mavproxy_chat/chat_openai.py

@@ -14,24 +14,54 @@
 import re
 from datetime import datetime
 from threading import Thread, Lock
+from typing_extensions import override
 import json
 import math
 from MAVProxy.modules.lib import param_help
 
 try:
-    from openai import OpenAI
+    from openai import OpenAI, AssistantEventHandler
 except Exception:
     print("chat: failed to import openai. See https://ardupilot.org/mavproxy/docs/modules/chat.html")
     exit()
 
 
+class EventHandler(AssistantEventHandler):
+    def __init__(self, chat_openai):
+        # record reference to chat_openai object
+        self.chat_openai = chat_openai
+
+        # initialise parent class (assistant event handler)
+        super().__init__()
+
+    @override
+    def on_event(self, event):
+        # record run id so that it can be cancelled if required
+        self.chat_openai.latest_run_id = event.data.id
+
+        # display the event in the status field
+        event_string_array = event.event.split(".")
+        event_string = event_string_array[-1]
+        self.chat_openai.send_status(event_string)
+
+        # requires_action events handled by function calls
+        if event.event == 'thread.run.requires_action':
+            self.chat_openai.handle_function_call(event)
+
+        # display reply text in the reply window
+        if (event.event == "thread.message.delta"):
+            stream_text = event.data.delta.content[0].text.value
+            self.chat_openai.send_reply(stream_text)
+
+
 class chat_openai():
-    def __init__(self, mpstate, status_cb=None, wait_for_command_ack_fn=None):
+    def __init__(self, mpstate, status_cb=None, reply_cb=None, wait_for_command_ack_fn=None):
         # keep reference to mpstate
         self.mpstate = mpstate
 
         # keep reference to status callback
         self.status_cb = status_cb
+        self.reply_cb = reply_cb
 
         # keep reference to wait_for_command_ack_fn
         self.wait_for_command_ack_fn = wait_for_command_ack_fn
@@ -48,7 +78,7 @@ def __init__(self, mpstate, status_cb=None, wait_for_command_ack_fn=None):
         self.client = None
         self.assistant = None
         self.assistant_thread = None
-        self.latest_run = None
+        self.latest_run_id = None
 
     # check connection to OpenAI assistant and connect if necessary
     # returns True if connection is good, False if not
@@ -102,22 +132,15 @@ def set_api_key(self, api_key_str):
 
     # cancel the active run
     def cancel_run(self):
-        # check the active thread and run
-        if (self.assistant_thread and self.latest_run is not None):
-            run_status = self.latest_run.status
-            if (run_status != "completed" and run_status != "cancelled" and
-                    run_status != "cancelling"):
-
-                # cancel the active run
-                self.client.beta.threads.runs.cancel(
-                    thread_id=self.assistant_thread.id,
-                    run_id=self.run.id
-                )
-            else:
-                if (self.latest_run.status == "completed"):
-                    print("Chat is completed, cannot be cancelled")
-                elif (self.latest_run.status == "cancelled"):
-                    print("Chat is cancelled")
+        # check the active thread and run id
+        if self.assistant_thread and self.latest_run_id is not None:
+            # cancel the run
+            self.client.beta.threads.runs.cancel(
+                thread_id=self.assistant_thread.id,
+                run_id=self.latest_run_id
+            )
+        else:
+            self.send_status("No active run to cancel")
 
     # send text to assistant
     def send_to_assistant(self, text):
@@ -126,7 +149,8 @@ def send_to_assistant(self, text):
 
             # check connection
             if not self.check_connection():
-                return "chat: failed to connect to OpenAI"
+                self.send_reply("chat: failed to connect to OpenAI")
+                return
 
             # create a new message
             input_message = self.client.beta.threads.messages.create(
@@ -135,91 +159,46 @@ def send_to_assistant(self, text):
                 content=text
             )
             if input_message is None:
-                return "chat: failed to create input message"
+                self.send_reply("chat: failed to create input message")
+                return
+
+            # create event handler
+            event_handler = EventHandler(self)
 
             # create a run
-            self.run = self.client.beta.threads.runs.create(
+            with self.client.beta.threads.runs.stream(
                 thread_id=self.assistant_thread.id,
-                assistant_id=self.assistant.id
-            )
-            if self.run is None:
-                return "chat: failed to create run"
-
-            # wait for run to complete
-            run_done = False
-            while not run_done:
-                # wait for one second
-                time.sleep(0.1)
-
-                # retrieve the run
-                self.latest_run = self.client.beta.threads.runs.retrieve(
+                assistant_id=self.assistant.id,
+                event_handler=event_handler
+            ) as stream:
+                stream.until_done()
+
+            # retrieve the run and print its final status
+            if self.assistant_thread and self.latest_run_id:
+                run = self.client.beta.threads.runs.retrieve(
                     thread_id=self.assistant_thread.id,
-                    run_id=self.run.id
+                    run_id=self.latest_run_id
                 )
-
-                # init failure message
-                failure_message = None
-
-                # check run status
-                if self.latest_run.status in ["queued", "in_progress", "cancelling"]:
-                    run_done = False
-                elif self.latest_run.status in ["cancelled", "completed", "expired"]:
-                    run_done = True
-                elif self.latest_run.status in ["failed"]:
-                    failure_message = self.latest_run.last_error.message
-                    run_done = True
-                elif self.latest_run.status in ["requires_action"]:
-                    self.handle_function_call(self.latest_run)
-                    run_done = False
-                else:
-                    print("chat: unrecognised run status" + self.latest_run.status)
-                    run_done = True
-
-                # send status to status callback
-                status_message = self.latest_run.status
-                if failure_message is not None:
-                    status_message = status_message + ": " + failure_message
-                self.send_status(status_message)
-
-            # retrieve messages on the thread
-            reply_messages = self.client.beta.threads.messages.list(self.assistant_thread.id,
-                                                                    order="asc",
-                                                                    after=input_message.id)
-
-            if (self.latest_run.status == "cancelled"):
-                return "cancelled successfully"
-            elif reply_messages is None:
-                return "chat: failed to retrieve messages"
-
-            # concatenate all messages into a single reply skipping the first which is our question
-            reply = ""
-            need_newline = False
-            for message in reply_messages.data:
-                reply = reply + message.content[0].text.value
-                if need_newline:
-                    reply = reply + "\n"
-                need_newline = True
-
-            if reply is None or reply == "":
-                return "chat: failed to retrieve latest reply"
-            return reply
+                if run is not None:
+                    self.send_status(run.status)
+            else:
+                self.send_status("done")
 
     # handle function call request from assistant
-    # on success this returns the text response that should be sent to the assistant, returns None on failure
-    def handle_function_call(self, run):
+    def handle_function_call(self, event):
 
         # sanity check required action (this should never happen)
-        if run.required_action is None:
+        if (event.event != "thread.run.requires_action"):
             print("chat::handle_function_call: assistant function call empty")
-            return None
+            return
 
         # check format
-        if run.required_action.submit_tool_outputs is None:
+        if event.data.required_action.submit_tool_outputs is None:
             print("chat::handle_function_call: submit tools outputs empty")
-            return None
+            return
 
         tool_outputs = []
-        for tool_call in run.required_action.submit_tool_outputs.tool_calls:
+        for tool_call in event.data.required_action.submit_tool_outputs.tool_calls:
             # init output to None
             output = "invalid function call"
 
@@ -274,13 +253,26 @@ def handle_function_call(self, run):
 
         # send function replies to assistant
         try:
-            self.client.beta.threads.runs.submit_tool_outputs(
-                thread_id=run.thread_id,
-                run_id=run.id,
-                tool_outputs=tool_outputs)
+            stream = self.client.beta.threads.runs.submit_tool_outputs(
+                thread_id=event.data.thread_id,
+                run_id=event.data.id,
+                tool_outputs=tool_outputs,
+                stream=True)
+
+            for event in stream:
+                # requires_action events handled by function calls
+                if event.event == 'thread.run.requires_action':
+                    self.handle_function_call(event)
+
+                # display reply text in the reply window
+                if (event.event == "thread.message.delta"):
+                    stream_text = event.data.delta.content[0].text.value
+                    self.send_reply(stream_text)
+
         except Exception:
             print("chat: error replying to function call")
             print(tool_outputs)
+            return
 
     # get the current date and time in the format, Saturday, June 24, 2023 6:14:14 PM
     def get_current_datetime(self, arguments):
@@ -800,6 +792,10 @@ def send_status(self, status):
         if self.status_cb is not None:
             self.status_cb(status)
 
+    def send_reply(self, reply):
+        if self.reply_cb is not None:
+            self.reply_cb(reply)
+
     # returns true if string contains regex characters
     def contains_regex(self, string):
         regex_characters = ".^$*+?{}[]\\|()"

MAVProxy/modules/mavproxy_chat/chat_window.py

@@ -18,7 +18,8 @@ def __init__(self, mpstate, wait_for_command_ack_fn):
         self.mpstate = mpstate
 
         # create chat_openai object
-        self.chat_openai = chat_openai.chat_openai(self.mpstate, self.set_status_text, wait_for_command_ack_fn)
+        self.chat_openai = chat_openai.chat_openai(self.mpstate, self.set_status_text, self.append_chat_replies,
+                                                   wait_for_command_ack_fn)
 
         # create chat_voice_to_text object
         self.chat_voice_to_text = chat_voice_to_text.chat_voice_to_text()
@@ -182,15 +183,11 @@ def send_text_to_assistant(self):
         wx.CallAfter(self.text_input.Clear)
 
         # copy user input text to reply box
-        orig_text_attr = self.text_reply.GetDefaultStyle()
         wx.CallAfter(self.text_reply.SetDefaultStyle, wx.TextAttr(wx.RED))
-        wx.CallAfter(self.text_reply.AppendText, send_text + "\n")
+        wx.CallAfter(self.text_reply.AppendText, "\n" + send_text + "\n")
 
-        # send text to assistant and place reply in reply box
-        reply = self.chat_openai.send_to_assistant(send_text)
-        if reply:
-            wx.CallAfter(self.text_reply.SetDefaultStyle, orig_text_attr)
-            wx.CallAfter(self.text_reply.AppendText, reply + "\n\n")
+        # send text to assistant. replies will be handled by append_chat_replies
+        self.chat_openai.send_to_assistant(send_text)
 
         # reenable buttons and text input (can't be done from a thread or must use CallAfter)
         # disable the cancel button
@@ -206,3 +203,8 @@ def send_text_to_assistant(self):
     # set status text
     def set_status_text(self, text):
         wx.CallAfter(self.text_status.SetValue, text)
+
+    # append chat to reply box
+    def append_chat_replies(self, text):
+        wx.CallAfter(self.text_reply.SetDefaultStyle, wx.TextAttr(wx.BLACK))
+        wx.CallAfter(self.text_reply.AppendText, text)