Better Python API, better KeyboardInturrupt

Author: KillianLucas

interpreter_1/cli.py

@@ -226,11 +226,17 @@ def load_interpreter_thread(args):
         spinner.stop()
     print()
     global_interpreter.messages = [{"role": "user", "content": message}]
-    async for _ in global_interpreter.async_respond():
-        pass
+    try:
+        async for _ in global_interpreter.async_respond():
+            pass
+    except KeyboardInterrupt:
+        global_interpreter._spinner.stop()
+    except asyncio.CancelledError:
+        global_interpreter._spinner.stop()
     print()
 
-    return global_interpreter
+    if global_interpreter.interactive:
+        await global_interpreter.async_chat()
 
 
 def parse_args():
@@ -297,17 +303,10 @@ def main():
             global_interpreter.server()
             return
 
-        # Run async portion
-        interpreter = asyncio.run(async_main(args))
-        # If we got an interpreter back and it's interactive, start chat in sync context
-        if interpreter and interpreter.interactive:
-            interpreter.chat()
-
+        asyncio.run(async_main(args))
     except KeyboardInterrupt:
-        print("KeyboardInterrupt")
         sys.exit(0)
     except asyncio.CancelledError:
-        print("CancelledError")
         sys.exit(0)
 
 

interpreter_1/interpreter.py

@@ -12,7 +12,7 @@
 
 from readchar import readchar
 
-from .misc.get_input import get_input
+from .misc.get_input import async_get_input
 
 # Third-party imports
 os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
@@ -227,11 +227,14 @@ def default_system_message(self):
 
         return system_message
 
-    async def async_respond(self):
+    async def async_respond(self, user_input=None):
         """
         Agentic sampling loop for the assistant/tool interaction.
         Yields chunks and maintains message history on the interpreter instance.
         """
+        if user_input:
+            self.messages.append({"role": "user", "content": user_input})
+
         tools = []
         if "interpreter" in self.tools:
             tools.append(BashTool())
@@ -925,14 +928,26 @@ def _handle_command(self, cmd: str, parts: list[str]) -> bool:
         return self._command_handler.handle_command(cmd, parts)
 
     def chat(self):
-        """
-        Interactive mode
-        """
+        """Chat with the interpreter. Handles both sync and async contexts."""
+        try:
+            loop = asyncio.get_running_loop()
+            # If we get here, there is a running event loop
+            loop.create_task(self.async_chat())
+        except RuntimeError:
+            # No running event loop, create one
+            asyncio.run(self.async_chat())
+
+    async def async_chat(self):
+        original_message_length = len(self.messages)
+
         try:
             message_count = 0
             while True:
-                user_input = get_input()
-                print("")
+                try:
+                    user_input = await async_get_input()
+                except KeyboardInterrupt:
+                    print()
+                    return self.messages[original_message_length:]
 
                 message_count += 1  # Increment counter after each message
 
@@ -943,23 +958,23 @@ def chat(self):
                         continue
 
                 if user_input == "":
-                    if message_count in range(4, 7):
+                    if message_count in range(8, 11):
                         print("Error: Cat is asleep on Enter key\n")
                     else:
                         print("Error: No input provided\n")
                     continue
 
-                self.messages.append({"role": "user", "content": user_input})
-
-                for _ in self.respond():
-                    pass
+                try:
+                    print()
+                    async for _ in self.async_respond(user_input):
+                        pass
+                except KeyboardInterrupt:
+                    self._spinner.stop()
+                except asyncio.CancelledError:
+                    self._spinner.stop()
 
                 print()
-        except KeyboardInterrupt:
-            self._spinner.stop()
-            print()
-            pass
-        except Exception as e:
+        except:
             self._spinner.stop()
             print(traceback.format_exc())
             print("\n\n\033[91mAn error has occurred.\033[0m")
@@ -976,35 +991,34 @@ def chat(self):
                     self._report_error("".join(traceback.format_exc()))
             exit(1)
 
-    async def _consume_generator(self, generator):
-        """Consume the async generator from async_respond"""
-        async for chunk in generator:
-            yield chunk
+    def respond(self, user_input=None, stream=False):
+        """Sync method to respond to user input if provided, or to the messages in self.messages."""
+        if user_input:
+            self.messages.append({"role": "user", "content": user_input})
 
-    def respond(self):
-        """
-        Synchronous wrapper around async_respond.
-        Yields chunks from the async generator.
-        """
+        if stream:
+            return self._sync_respond_stream()
+        else:
+            original_message_length = len(self.messages)
+            for _ in self._sync_respond_stream():
+                pass
+            return self.messages[original_message_length:]
+
+    def _sync_respond_stream(self):
+        """Synchronous generator that yields responses. Only use in synchronous contexts."""
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
         try:
-            loop = asyncio.get_event_loop()
-        except RuntimeError:
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-
-        async def run():
-            async for chunk in self.async_respond():
-                yield chunk
-
-        agen = run()
-        while True:
-            try:
-                chunk = loop.run_until_complete(anext(agen))
-                yield chunk
-            except StopAsyncIteration:
-                break
-
-        return self.messages
+            # Convert async generator to sync generator
+            async_gen = self.async_respond()
+            while True:
+                try:
+                    chunk = loop.run_until_complete(async_gen.__anext__())
+                    yield chunk
+                except StopAsyncIteration:
+                    break
+        finally:
+            loop.close()
 
     def server(self):
         """

interpreter_1/misc/get_input.py

@@ -51,7 +51,6 @@ def _(event):
         include_default_pygments_style=False,
         input_processors=[],
         enable_system_prompt=False,
-        wrap_lines=False,
     )
     return result
 
@@ -101,6 +100,5 @@ def _(event):
         include_default_pygments_style=False,
         input_processors=[],
         enable_system_prompt=False,
-        wrap_lines=False,
     )
     return result

interpreter_1/server.py

@@ -56,25 +56,7 @@ async def chat_completion(self, request: Request):
                 self._stream_response(), media_type="text/event-stream"
             )
 
-        # For non-streaming, collect all chunks
-        response_text = ""
-        for chunk in self.interpreter.respond():
-            if chunk.get("type") == "chunk":
-                response_text += chunk["chunk"]
-
-        return {
-            "id": "chatcmpl-" + str(time.time()),
-            "object": "chat.completion",
-            "created": int(time.time()),
-            "model": req.model or self.interpreter.model,
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {"role": "assistant", "content": response_text},
-                    "finish_reason": "stop",
-                }
-            ],
-        }
+        raise NotImplementedError("Non-streaming is not supported yet")
 
     async def _stream_response(self):
         """Stream the response in OpenAI-compatible format"""

interpreter_1/tools/bash.py

@@ -1,6 +1,7 @@
 import asyncio
 import os
 import shutil
+import traceback
 from typing import ClassVar, Literal
 
 import pyte
@@ -91,9 +92,17 @@ async def run(self, command: str):
             final_output = self._get_screen_text()
             return CLIResult(output=final_output if final_output else "<No output>")
 
-        except (KeyboardInterrupt, asyncio.CancelledError):
+        except KeyboardInterrupt as e:
             self.stop()
             return CLIResult(output="Command cancelled by user.")
+        except asyncio.CancelledError as e:
+            self.stop()
+            return CLIResult(output="Command cancelled by user.")
+        except Exception as e:
+            print("Unexpected error")
+            traceback.print_exc()
+            self.stop()
+            return CLIResult(output=f"Command failed with error: {e}")
 
 
 class BashTool(BaseAnthropicTool):