launch container from the tool init

examples/web_agent/demo.sh

@@ -1,31 +1,12 @@
 #!/bin/bash
-
-# Add trap for cleanup without stopping podman machine
-SCRIPT_DIR=$(dirname "$0")
-cleanup() {
-    echo "Cleaning up..."
-    pkill -f "$SCRIPT_DIR/http_server.py"
-    pkill -f "$SCRIPT_DIR/chat.py"
-    exit 0
-}
-
-pkill -f "http_server.py"
-pkill -f "chat.py"
-
-# Set up trap for SIGINT and SIGTERM
-trap cleanup SIGINT SIGTERM
-
 if [[ "$(uname)" != "Darwin" ]]; then
     echo "Error: This script only works on macOS"
     exit 1
 fi
 
 # run podman and container
-podman machine set --user-mode-networking
-export DOCKER_HOST=http+unix://$(podman machine inspect --format '{{.ConnectionInfo.PodmanSocket.Path}}')
 if ! command -v podman &> /dev/null; then
     echo "Podman is not installed, installing..."
-
     if ! command -v brew &> /dev/null; then
         echo "Error: Homebrew is not installed. Please install it first."
         echo "Visit https://brew.sh for installation instructions."
@@ -45,54 +26,37 @@ if ! podman machine list | grep -q "Currently running"; then
         exit 1
     fi
 fi
-if ! podman ps --format "{{.Names}}" | grep -q "^computer$"; then
-    echo "No computer container found, starting one"
-    if ! podman images computer | grep -q "computer"; then
-        echo "No computer image found, building one"
-        podman images
-        podman build -t computer:latest tapeagents/tools/computer/
-        if [ $? -ne 0 ]; then
-            echo "Failed to build computer image"
-            exit 1
-        fi
-    else
-        echo "Computer image found"
+podman machine set --user-mode-networking
+export DOCKER_HOST=http+unix://$(podman machine inspect --format '{{.ConnectionInfo.PodmanSocket.Path}}')
+if ! podman images computer | grep -q "computer"; then
+    echo "No computer image found, building one"
+    podman images
+    podman build -t computer:latest tapeagents/tools/computer/
+    if [ $? -ne 0 ]; then
+        echo "Failed to build computer image"
+        exit 1
     fi
-    # Initialize the log file
-    > /tmp/demo_stdout.log
-    echo "Starting computer container"
-    ./tapeagents/tools/computer/run.sh >> /tmp/demo_stdout.log 2>&1 &
-    echo -n "Waiting for computer container to start"
-    # Wait up to 30 seconds for computer container to be running
-    for i in {1..30}; do
-        if podman ps | grep -q "computer"; then
-            break
-        fi
-        sleep 1
-        echo -n "."
-        if [ $i -eq 30 ]; then
-            echo "Error: Computer container failed to start within 30 seconds"
-            exit 1
-        fi
-    done
-    echo "."
-    echo -n "Waiting for API init"
-    while ! grep -q "Uvicorn running on" /tmp/demo_stdout.log; do
-        sleep 1
-        echo -n "."
-    done
-    echo "."
-    echo "Computer started"
 fi
 
 # install dependencies
 uv sync --all-extras
 
-# start scripts
-echo "Starting Chat UI"
+SCRIPT_DIR=$(dirname "$0")
+
+# Add trap for cleanup
+cleanup() {
+    echo "Cleaning up..."
+    pkill -f "$SCRIPT_DIR/chat.py"
+    exit 0
+}
+trap cleanup SIGINT SIGTERM
+
 mkdir -p .cache
-uv run $SCRIPT_DIR/http_server.py >> /tmp/demo_stdout.log 2>&1 &
-STREAMLIT_SERVER_PORT=8501 uv run -m streamlit run $SCRIPT_DIR/chat.py --server.headless true >> /tmp/demo_stdout.log 2>&1 &
+# Initialize the log file
+> /tmp/demo_stdout.log
+uv run -m tapeagents.tools.computer
+uv run -m http.server 8080 --directory $SCRIPT_DIR/static_content >> /tmp/demo_stdout.log 2>&1 &
+uv run -m streamlit run $SCRIPT_DIR/chat.py --server.headless true >> /tmp/demo_stdout.log 2>&1 &
 sleep 3
 echo "Tapeagents Operator is ready"
 echo "Open http://localhost:8080 in your browser to begin"

tapeagents/tools/computer/Dockerfile

@@ -15,5 +15,6 @@ RUN python3 -m pip install pyautogui uvicorn fastapi python-multipart pillow pys
 RUN apt -y install xdotool
 
 COPY api.py /api.py
+COPY start.sh /custom-cont-init.d/start.sh
 RUN ln -s /usr/bin/python3 /usr/bin/python
 RUN ln -s /config /home/user

tapeagents/tools/computer/__init__.py

@@ -1,3 +1,68 @@
 """
 Computer tool, based on claude computer use demo
 """
+
+import atexit
+import logging
+import os
+import time
+
+import podman
+
+logger = logging.getLogger("computer")
+logger.setLevel(logging.INFO)
+formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(name)s:%(lineno)d - %(funcName)s - %(message)s")
+console_handler = logging.StreamHandler()
+console_handler.setFormatter(formatter)
+logger.addHandler(console_handler)
+
+
+def launch_container(container_image="computer", container_name="computer", stop_at_exit=False):
+    podman_client = podman.from_env()
+    try:
+        container = podman_client.containers.get(container_name)
+    except podman.errors.NotFound:
+        logger.info(f"Creating container from image '{container_image}'")
+        home_dir = os.path.join(os.getcwd(), ".computer")
+        zip_home = os.path.join(os.path.dirname(__file__), "home.zip")
+        if os.path.exists(home_dir):
+            os.system(f"rm -rf {home_dir}")
+            logger.info(f"Removed existing home directory: {home_dir}")
+        else:
+            os.makedirs(home_dir)
+        os.system(f"unzip -qq -o {zip_home} -d {home_dir}")
+        logger.info(f"Recreated home directory from zip: {zip_home}")
+        container = podman_client.containers.create(
+            container_image,
+            name=container_name,
+            auto_remove=True,
+            ports={
+                "3000": 3000,
+                "8000": 8000,
+            },
+            mounts=[
+                {
+                    "type": "bind",
+                    "source": os.path.join(home_dir, "home"),
+                    "target": "/config",
+                }
+            ],
+        )
+        container.start()
+        logger.info("Starting..")
+    while container.status != "running":
+        time.sleep(0.2)
+        container.reload()
+    logger.info("Container ready")
+
+    if stop_at_exit:
+
+        def cleanup() -> None:
+            try:
+                logger.info(f"Stopping container {container.name}")
+                container.stop()
+            except podman.errors.NotFound:
+                pass
+            atexit.unregister(cleanup)
+
+        atexit.register(cleanup)

tapeagents/tools/computer/__main__.py

@@ -0,0 +1,4 @@
+from tapeagents.tools.computer import launch_container
+
+if __name__ == "__main__":
+    launch_container()

tapeagents/tools/computer/api.py

@@ -11,9 +11,9 @@
 from fastapi import FastAPI, HTTPException
 from PIL import Image
 
-logger = logging.getLogger("API")
+logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
-formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(lineno)d - %(funcName)s - %(message)s")
+formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(name)s:%(lineno)d - %(funcName)s - %(message)s")
 console_handler = logging.StreamHandler()
 console_handler.setFormatter(formatter)
 logger.addHandler(console_handler)

tapeagents/tools/computer/remote.py

@@ -11,9 +11,8 @@
 from tapeagents.core import Action
 from tapeagents.steps import ImageObservation
 from tapeagents.tools.base import Multitool
-from tapeagents.tools.grounding import GroundingModel
-
-from .steps import (
+from tapeagents.tools.computer import launch_container
+from tapeagents.tools.computer.steps import (
     ComputerObservation,
     GetCursorPositionAction,
     KeyPressAction,
@@ -23,10 +22,11 @@
     RunTerminalCommand,
     TypeTextAction,
 )
+from tapeagents.tools.grounding import GroundingModel
 
-logger = logging.getLogger("remote")
+logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
-formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(lineno)d - %(funcName)s - %(message)s")
+formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(name)s:%(lineno)d - %(funcName)s - %(message)s")
 console_handler = logging.StreamHandler()
 console_handler.setFormatter(formatter)
 logger.addHandler(console_handler)
@@ -80,6 +80,8 @@ class RemoteComputer(Multitool):
     computer_url: str = Field(description="Remote tool API URL")
     use_grounding: bool = Field(description="Whether to use grounding model", default=True)
     grounding_api_url: str = Field(description="Grounding API URL")
+    container_image: str = "computer"
+    container_name: str = "computer"
 
     def model_post_init(self, __context):
         self._grounding = GroundingModel(url=self.grounding_api_url)
@@ -101,6 +103,7 @@ def model_post_init(self, __context):
             self._action_map[MouseXYMoveAction] = self.remote_execute_action
             self._action_map[GetCursorPositionAction] = self.remote_execute_action
         self.actions = tuple(self._action_map.keys())
+        launch_container(self.container_image, self.container_name, stop_at_exit=True)
 
     def execute_action(self, action: Action) -> ImageObservation:
         action_type = type(action)
@@ -151,3 +154,6 @@ def get_screen(self) -> Image:
 
     def reset(self):
         self.remote_execute_action(RunTerminalCommand(command='xdotool search "" windowkill %@'))
+
+    def close(self):
+        return super().close()

tapeagents/tools/computer/start.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+HOME=/config DISPLAY=:1 python3 /api.py > /tmp/api.log 2>&1 &
+echo "API started"