Async Interpreter

Author: KillianLucas

interpreter/__init__.py

@@ -1,3 +1,4 @@
+from .core.async_core import AsyncInterpreter
 from .core.computer.terminal.base_language import BaseLanguage
 from .core.core import OpenInterpreter
 

interpreter/core/archived_server.py renamed to interpreter/core/archived_server_1.py

@@ -1,35 +1,190 @@
+import asyncio
 import json
 import threading
-import time
+import traceback
+from typing import Any, Dict
 
-from core import OpenInterpreter
+from .core import OpenInterpreter
 
+try:
+    import janus
+    import uvicorn
+    from fastapi import APIRouter, FastAPI, WebSocket
+except:
+    # Server dependencies are not required by the main package.
+    pass
 
-class AsyncOpenInterpreter(OpenInterpreter):
+
+class AsyncInterpreter(OpenInterpreter):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.async_thread = None
-        self.input_queue
-        self.output_queue
+
+        self.respond_thread = None
+        self.stop_event = threading.Event()
+        self.output_queue = None
+
+        self.server = Server(self)
 
     async def input(self, chunk):
         """
-        Expects a chunk in streaming LMC format.
+        Accumulates LMC chunks onto interpreter.messages.
+        When it hits an "end" flag, calls interpreter.respond().
         """
-        try:
-            chunk = json.loads(chunk)
-        except:
-            pass
 
         if "start" in chunk:
-            self.async_thread.join()
+            # If the user is starting something, the interpreter should stop.
+            if self.respond_thread is not None and self.respond_thread.is_alive():
+                self.stop_event.set()
+                self.respond_thread.join()
+            self.accumulate(chunk)
+        elif "content" in chunk:
+            self.accumulate(chunk)
         elif "end" in chunk:
-            if self.async_thread is None or not self.async_thread.is_alive():
-                self.async_thread = threading.Thread(target=self.complete)
-            self.async_thread.start()
-        else:
-            await self._add_to_queue(self._input_queue, chunk)
-
-    async def output(self, *args, **kwargs):
-        # Your async output code here
-        pass
+            # If the user is done talking, the interpreter should respond.
+            self.stop_event.clear()
+            print("Responding.")
+            self.respond_thread = threading.Thread(target=self.respond)
+            self.respond_thread.start()
+
+    async def output(self):
+        if self.output_queue == None:
+            self.output_queue = janus.Queue()
+        return await self.output_queue.async_q.get()
+
+    def respond(self):
+        for chunk in self._respond_and_store():
+            print(chunk.get("content", ""), end="")
+            if self.stop_event.is_set():
+                return
+            self.output_queue.sync_q.put(chunk)
+
+        self.output_queue.sync_q.put(
+            {"role": "server", "type": "status", "content": "complete"}
+        )
+
+    def accumulate(self, chunk):
+        """
+        Accumulates LMC chunks onto interpreter.messages.
+        """
+        if type(chunk) == dict:
+            if chunk.get("format") == "active_line":
+                # We don't do anything with these.
+                pass
+
+            elif "start" in chunk:
+                chunk_copy = (
+                    chunk.copy()
+                )  # So we don't modify the original chunk, which feels wrong.
+                chunk_copy.pop("start")
+                chunk_copy["content"] = ""
+                self.messages.append(chunk_copy)
+
+            elif "content" in chunk:
+                self.messages[-1]["content"] += chunk["content"]
+
+        elif type(chunk) == bytes:
+            if self.messages[-1]["content"] == "":  # We initialize as an empty string ^
+                self.messages[-1]["content"] = b""  # But it actually should be bytes
+            self.messages[-1]["content"] += chunk
+
+
+def create_router(async_interpreter):
+    router = APIRouter()
+
+    @router.get("/heartbeat")
+    async def heartbeat():
+        return {"status": "alive"}
+
+    @router.websocket("/")
+    async def websocket_endpoint(websocket: WebSocket):
+        await websocket.accept()
+        try:
+
+            async def receive_input():
+                while True:
+                    try:
+                        data = await websocket.receive()
+
+                        if data.get("type") == "websocket.receive" and "text" in data:
+                            data = json.loads(data["text"])
+                            await async_interpreter.input(data)
+                        elif (
+                            data.get("type") == "websocket.disconnect"
+                            and data.get("code") == 1000
+                        ):
+                            print("Disconnecting.")
+                            return
+                        else:
+                            print("Invalid data:", data)
+                            continue
+
+                    except Exception as e:
+                        error_message = {
+                            "role": "server",
+                            "type": "error",
+                            "content": traceback.format_exc() + "\n" + str(e),
+                        }
+                        await websocket.send_text(json.dumps(error_message))
+
+            async def send_output():
+                while True:
+                    try:
+                        output = await async_interpreter.output()
+
+                        if isinstance(output, bytes):
+                            await websocket.send_bytes(output)
+                        else:
+                            await websocket.send_text(json.dumps(output))
+                    except Exception as e:
+                        traceback.print_exc()
+                        error_message = {
+                            "role": "server",
+                            "type": "error",
+                            "content": traceback.format_exc() + "\n" + str(e),
+                        }
+                        await websocket.send_text(json.dumps(error_message))
+
+            await asyncio.gather(receive_input(), send_output())
+        except Exception as e:
+            traceback.print_exc()
+            try:
+                error_message = {
+                    "role": "server",
+                    "type": "error",
+                    "content": traceback.format_exc() + "\n" + str(e),
+                }
+                await websocket.send_text(json.dumps(error_message))
+            except:
+                # If we can't send it, that's fine.
+                pass
+        finally:
+            await websocket.close()
+
+    @router.post("/settings")
+    async def settings(payload: Dict[str, Any]):
+        for key, value in payload.items():
+            print(f"Updating settings: {key} = {value}")
+            if key in ["llm", "computer"] and isinstance(value, dict):
+                for sub_key, sub_value in value.items():
+                    setattr(getattr(async_interpreter, key), sub_key, sub_value)
+            else:
+                setattr(async_interpreter, key, value)
+
+        return {"status": "success"}
+
+    return router
+
+
+class Server:
+    def __init__(self, async_interpreter, host="0.0.0.0", port=8000):
+        self.app = FastAPI()
+        router = create_router(async_interpreter)
+        self.app.include_router(router)
+        self.host = host
+        self.port = port
+        self.uvicorn_server = uvicorn.Server(
+            config=uvicorn.Config(app=self.app, host=self.host, port=self.port)
+        )
+
+    def run(self):
+        uvicorn.run(self.app, host=self.host, port=self.port)

interpreter/core/server.py renamed to interpreter/core/archived_server_2.py

@@ -236,6 +236,13 @@ def detect_active_line(self, line):
 
     def _capture_output(self, message_queue):
         while True:
+            # For async usage
+            if (
+                hasattr(self.computer.interpreter, "stop_event")
+                and self.computer.interpreter.stop_event.is_set()
+            ):
+                break
+
             if self.listener_thread:
                 try:
                     output = message_queue.get(timeout=0.1)

interpreter/core/async_core.py

@@ -20,12 +20,6 @@
 from .utils.telemetry import send_telemetry
 from .utils.truncate_output import truncate_output
 
-try:
-    from .server import server
-except:
-    # Dependencies for server are not generally required
-    pass
-
 
 class OpenInterpreter:
     """
@@ -141,14 +135,6 @@ def __init__(
         self.empty_code_output_template = empty_code_output_template
         self.code_output_sender = code_output_sender
 
-    def server(self, *args, **kwargs):
-        try:
-            server(self)
-        except ImportError:
-            display_markdown_message(
-                "Missing dependencies for the server, please run `pip install open-interpreter[server]` and try again."
-            )
-
     def local_setup(self):
         """
         Opens a wizard that lets terminal users pick a local model.
@@ -313,6 +299,7 @@ def _respond_and_store(self):
         Pulls from the respond stream, adding delimiters. Some things, like active_line, console, confirmation... these act specially.
         Also assembles new messages and adds them to `self.messages`.
         """
+        self.verbose = False
 
         # Utility function
         def is_active_line_chunk(chunk):
@@ -321,6 +308,10 @@ def is_active_line_chunk(chunk):
         last_flag_base = None
 
         for chunk in respond(self):
+            # For async usage
+            if hasattr(self, "stop_event") and self.stop_event.is_set():
+                break
+
             if chunk["content"] == "":
                 continue
 
@@ -330,7 +321,10 @@ def is_active_line_chunk(chunk):
                 if last_flag_base:
                     yield {**last_flag_base, "end": True}
                     last_flag_base = None
-                yield chunk
+
+                if self.auto_run == False:
+                    yield chunk
+
                 # We want to append this now, so even if content is never filled, we know that the execution didn't produce output.
                 # ... rethink this though.
                 self.messages.append(

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

@@ -4,10 +4,12 @@
 import litellm
 
 litellm.suppress_debug_info = True
+import json
 import subprocess
 import time
 import uuid
 
+import requests
 import tokentrim as tt
 
 from ...terminal_interface.utils.display_markdown_message import (
@@ -289,8 +291,6 @@ def load(self):
             return
 
         if self.model.startswith("ollama/"):
-            # WOAH we should also hit up ollama and set max_tokens and context_window based on the LLM. I think they let u do that
-
             model_name = self.model.replace("ollama/", "")
             try:
                 # List out all downloaded ollama models. Will fail if ollama isn't installed
@@ -315,20 +315,36 @@ def load(self):
                 self.interpreter.display_message(f"\nDownloading {model_name}...\n")
                 subprocess.run(["ollama", "pull", model_name], check=True)
 
+            # Get context window if not set
+            if self.context_window == None:
+                response = requests.post(
+                    "http://localhost:11434/api/show", json={"name": model_name}
+                )
+                model_info = response.json().get("model_info", {})
+                context_length = None
+                for key in model_info:
+                    if "context_length" in key:
+                        context_length = model_info[key]
+                        break
+                if context_length is not None:
+                    self.context_window = context_length
+            if self.max_tokens == None:
+                if self.context_window != None:
+                    self.max_tokens = int(self.context_window * 0.8)
+
             # Send a ping, which will actually load the model
-            # print(f"\nLoading {model_name}...\n")
+            print(f"Loading {model_name}...\n")
 
             old_max_tokens = self.max_tokens
             self.max_tokens = 1
             self.interpreter.computer.ai.chat("ping")
             self.max_tokens = old_max_tokens
 
-            # self.interpreter.display_message("\n*Model loaded.*\n")
+            self.interpreter.display_message("*Model loaded.*\n")
 
         # Validate LLM should be moved here!!
 
         self._is_loaded = True
-        return
 
 
 def fixed_litellm_completions(**params):

interpreter/core/core.py

@@ -27,8 +27,3 @@
 # Misc settings
 interpreter.auto_run = False
 interpreter.offline = True
-
-# Final message
-interpreter.display_message(
-    f"> Model set to `{interpreter.llm.model}`\n\n**Open Interpreter** will require approval before running code.\n\nUse `interpreter -y` to bypass this.\n\nPress `CTRL-C` to exit.\n"
-)