Merge pull request #26 from EnvisionMindCa/codex/make-vm-asynchronous-and-modify-llm-flow

README.md

@@ -8,8 +8,9 @@ conversations can be resumed with context. One example tool is included:
 * **execute_terminal** – Executes a shell command inside a persistent Linux VM
   with network access. Use it to read uploaded documents under ``/data`` or run
   other commands. Output from ``stdout`` and ``stderr`` is captured and
-  returned. The VM is created when a chat session starts and reused for all
-  subsequent tool calls.
+  returned. Commands run asynchronously so the assistant can continue
+  responding while they execute. The VM is created when a chat session starts
+  and reused for all subsequent tool calls.
 
 The application injects a robust system prompt on each request. The prompt
 guides the model to plan tool usage, execute commands sequentially and

src/__init__.py

@@ -1,5 +1,11 @@
 from .chat import ChatSession
-from .tools import execute_terminal, set_vm
+from .tools import execute_terminal, execute_terminal_async, set_vm
 from .vm import LinuxVM
 
-__all__ = ["ChatSession", "execute_terminal", "set_vm", "LinuxVM"]
+__all__ = [
+    "ChatSession",
+    "execute_terminal",
+    "execute_terminal_async",
+    "set_vm",
+    "LinuxVM",
+]

src/chat.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from typing import List
 import json
+import asyncio
 import shutil
 from pathlib import Path
 
@@ -26,7 +27,7 @@ from .db import (
 )
 from .log import get_logger
 from .schema import Msg
-from .tools import execute_terminal, set_vm
+from .tools import execute_terminal, execute_terminal_async, set_vm
 from .vm import VMRegistry
 
 _LOG = get_logger(__name__)
@@ -141,29 +142,77 @@ class ChatSession:
     ) -> ChatResponse:
         while depth < MAX_TOOL_CALL_DEPTH and response.message.tool_calls:
             for call in response.message.tool_calls:
-                if call.function.name == "execute_terminal":
-                    result = execute_terminal(**call.function.arguments)
-                else:
+                if call.function.name != "execute_terminal":
                     _LOG.warning("Unsupported tool call: %s", call.function.name)
                     result = f"Unsupported tool: {call.function.name}"
+                    messages.append(
+                        {
+                            "role": "tool",
+                            "name": call.function.name,
+                            "content": result,
+                        }
+                    )
+                    DBMessage.create(
+                        conversation=conversation,
+                        role="tool",
+                        content=result,
+                    )
+                    continue
 
-                messages.append(
-                    {
-                        "role": "tool",
-                        "name": call.function.name,
-                        "content": str(result),
-                    }
+                exec_task = asyncio.create_task(
+                    execute_terminal_async(**call.function.arguments)
                 )
-                DBMessage.create(
-                    conversation=conversation,
-                    role="tool",
-                    content=str(result),
+                follow_task = asyncio.create_task(self.ask(messages, think=True))
+
+                done, _ = await asyncio.wait(
+                    {exec_task, follow_task},
+                    return_when=asyncio.FIRST_COMPLETED,
                 )
 
-            nxt = await self.ask(messages, think=True)
-            self._store_assistant_message(conversation, nxt.message)
-            response = nxt
-            depth += 1
+                if exec_task in done:
+                    follow_task.cancel()
+                    try:
+                        await follow_task
+                    except Exception:
+                        pass
+                    result = await exec_task
+                    messages.append(
+                        {
+                            "role": "tool",
+                            "name": call.function.name,
+                            "content": result,
+                        }
+                    )
+                    DBMessage.create(
+                        conversation=conversation,
+                        role="tool",
+                        content=result,
+                    )
+                    nxt = await self.ask(messages, think=True)
+                    self._store_assistant_message(conversation, nxt.message)
+                    response = nxt
+                else:
+                    followup = await follow_task
+                    self._store_assistant_message(conversation, followup.message)
+                    messages.append(followup.message.model_dump())
+                    result = await exec_task
+                    messages.append(
+                        {
+                            "role": "tool",
+                            "name": call.function.name,
+                            "content": result,
+                        }
+                    )
+                    DBMessage.create(
+                        conversation=conversation,
+                        role="tool",
+                        content=result,
+                    )
+                    nxt = await self.ask(messages, think=True)
+                    self._store_assistant_message(conversation, nxt.message)
+                    response = nxt
+
+                depth += 1
 
         return response
 

src/config.py

@@ -13,10 +13,12 @@ UPLOAD_DIR: Final[str] = os.getenv("UPLOAD_DIR", str(Path.cwd() / "uploads"))
 SYSTEM_PROMPT: Final[str] = (
     "You are Starlette, a professional AI assistant with advanced tool orchestration. "
     "Always analyze the user's objective before responding. If tools are needed, "
-    "outline a step-by-step plan and invoke each tool sequentially, waiting for its "
-    "result before proceeding. Retry or adjust commands when they fail and summarize "
-    "important outputs to preserve context. Uploaded files live under /data and are "
-    "accessible via the execute_terminal tool. Continue using tools until you have "
-    "gathered everything required to produce an accurate answer, then craft a clear "
-    "and precise final response that fully addresses the request."
+    "outline a step-by-step plan and invoke each tool sequentially. Shell commands "
+    "execute asynchronously, so provide a brief interim reply while waiting. Once a "
+    "tool returns its result you will receive a tool message and must continue from "
+    "there. If the result arrives before your interim reply is complete, cancel the "
+    "reply and incorporate the tool output instead. Uploaded files live under /data "
+    "and are accessible via the execute_terminal tool. Continue using tools until "
+    "you have gathered everything required to produce an accurate answer, then craft "
+    "a clear and precise final response that fully addresses the request."
 )

src/tools.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
-__all__ = ["execute_terminal", "set_vm"]
+__all__ = ["execute_terminal", "execute_terminal_async", "set_vm"]
 
 import subprocess
 from typing import Optional
+import asyncio
 
 from .vm import LinuxVM
 
@@ -53,3 +54,9 @@ def execute_terminal(command: str) -> str:
     if completed.stderr:
         output = f"{output}\n{completed.stderr}" if output else completed.stderr
     return output.strip()
+
+
+async def execute_terminal_async(command: str) -> str:
+    """Asynchronously execute a shell command."""
+    loop = asyncio.get_running_loop()
+    return await loop.run_in_executor(None, execute_terminal, command)

src/vm.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import subprocess
+import asyncio
+from functools import partial
 import uuid
 from pathlib import Path
 
@@ -90,6 +92,12 @@ class LinuxVM:
             output = f"{output}\n{completed.stderr}" if output else completed.stderr
         return output.strip()
 
+    async def execute_async(self, command: str, *, timeout: int = 3) -> str:
+        """Asynchronously execute ``command`` inside the running VM."""
+        loop = asyncio.get_running_loop()
+        func = partial(self.execute, command, timeout=timeout)
+        return await loop.run_in_executor(None, func)
+
     def stop(self) -> None:
         """Terminate the VM if running."""
         if not self._running: