Added jupyter retries, minor robust fixes

Author: KillianLucas

interpreter/core/async_core.py

@@ -441,11 +441,16 @@ async def receive_input():
                             "type": "error",
                             "content": traceback.format_exc() + "\n" + str(e),
                         }
-                        await websocket.send_text(json.dumps(error_message))
-                        await websocket.send_text(json.dumps(complete_message))
-                        print("\n\n--- SENT ERROR: ---\n\n")
+                        if websocket.client_state == WebSocketState.CONNECTED:
+                            await websocket.send_text(json.dumps(error_message))
+                            await websocket.send_text(json.dumps(complete_message))
+                            print("\n\n--- SENT ERROR: ---\n\n")
+                        else:
+                            print(
+                                "\n\n--- ERROR (not sent due to disconnected state): ---\n\n"
+                            )
                         print(error)
-                        print("\n\n--- (ERROR ABOVE WAS SENT) ---\n\n")
+                        print("\n\n--- (ERROR ABOVE) ---\n\n")
 
             async def send_output():
                 while True:

interpreter/core/computer/terminal/languages/jupyter_language.py

@@ -5,10 +5,10 @@
 
 import ast
 import logging
-import sys
 import os
 import queue
 import re
+import sys
 import threading
 import time
 import traceback
@@ -21,8 +21,8 @@
 
 # When running from an executable, ipykernel calls itself infinitely
 # This is a workaround to detect it and launch it manually
-if 'ipykernel_launcher' in sys.argv:
-    if sys.path[0] == '':
+if "ipykernel_launcher" in sys.argv:
+    if sys.path[0] == "":
         del sys.path[0]
 
     from ipykernel import kernelapp as app
@@ -127,6 +127,7 @@ def run(self, code):
 
     def _execute_code(self, code, message_queue):
         def iopub_message_listener():
+            max_retries = 100
             while True:
                 # If self.finish_flag = True, and we didn't set it (we do below), we need to stop. That's our "stop"
                 if self.finish_flag == True:
@@ -138,6 +139,12 @@ def iopub_message_listener():
                     msg = self.kc.iopub_channel.get_msg(timeout=0.05)
                 except queue.Empty:
                     continue
+                except Exception as e:
+                    max_retries -= 1
+                    if max_retries < 0:
+                        raise
+                    print("Jupyter error, retrying:", str(e))
+                    continue
 
                 if DEBUG_MODE:
                     print("-----------" * 10)
@@ -248,6 +255,8 @@ def detect_active_line(self, line):
 
     def _capture_output(self, message_queue):
         while True:
+            time.sleep(0.1)
+
             # For async usage
             if (
                 hasattr(self.computer.interpreter, "stop_event")
@@ -263,10 +272,17 @@ def _capture_output(self, message_queue):
                     yield output
                 except queue.Empty:
                     if self.finish_flag:
-                        if DEBUG_MODE:
-                            print("we're done")
-                        break
-            time.sleep(0.1)
+                        time.sleep(0.1)
+
+                        try:
+                            output = message_queue.get(timeout=0.1)
+                            if DEBUG_MODE:
+                                print(output)
+                            yield output
+                        except queue.Empty:
+                            if DEBUG_MODE:
+                                print("we're done")
+                            break
 
     def stop(self):
         self.finish_flag = True

interpreter/core/llm/llm.py

@@ -286,7 +286,6 @@ def run(self, messages):
                     print(message)
                 print("\n")
             print("\n\n\n")
-            time.sleep(5)
 
         if self.supports_functions:
             # yield from run_function_calling_llm(self, params)

interpreter/core/llm/run_tool_calling_llm.py

@@ -312,9 +312,9 @@ def run_tool_calling_llm(llm, request_params):
             else:
                 # If name exists and it's not "execute" or "python" or "functions", who knows what's going on.
                 if "name" in accumulated_deltas["function_call"]:
-                    yield {
-                        "type": "code",
-                        "format": "python",
-                        "content": accumulated_deltas["function_call"]["name"],
-                    }
+                    # yield {
+                    #     "type": "code",
+                    #     "format": "python",
+                    #     "content": str(delta),
+                    # }
                     return

interpreter/terminal_interface/start_terminal_interface.py

@@ -22,6 +22,12 @@ def start_terminal_interface(interpreter):
     Meant to be used from the command line. Parses arguments, starts OI's terminal interface.
     """
 
+    # Instead use an async interpreter, which has a server. Set settings on that
+    if "--server" in sys.argv:
+        from interpreter import AsyncInterpreter
+
+        interpreter = AsyncInterpreter()
+
     arguments = [
         {
             "name": "profile",
@@ -371,12 +377,6 @@ def print_help(self, *args, **kwargs):
 
     args, unknown_args = parser.parse_known_args()
 
-    if args.server:
-        # Instead use an async interpreter, which has a server. Set settings on that
-        from interpreter import AsyncInterpreter
-
-        interpreter = AsyncInterpreter()
-
     # handle unknown arguments
     if unknown_args:
         print(f"\nUnrecognized argument(s): {unknown_args}")