agentscope-runtime 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

agentscope_runtime/sandbox/box/training_box/training_box.py

@@ -5,21 +5,13 @@ Module for the Training Sandbox implementation.
 This module provides a sandbox environment for training tasks
 with specific configuration and tool calling methods.
 """
-import platform
 from typing import Dict, Optional
 import os
 
+from ...utils import build_image_uri
 from ...registry import SandboxRegistry
 from ...enums import SandboxType
 from ...box.sandbox import Sandbox
-from ...constant import IMAGE_TAG
-
-
-def get_image_tag() -> str:
-    machine = platform.machine().lower()
-    if machine in ("arm64", "aarch64", "armv7l", "armv8"):
-        return f"{IMAGE_TAG}-arm64"
-    return IMAGE_TAG
 
 
 class TrainingSandbox(Sandbox):
@@ -214,7 +206,7 @@ class TrainingSandbox(Sandbox):
 
 
 @SandboxRegistry.register(
-    f"agentscope/runtime-sandbox-appworld:{get_image_tag()}",
+    build_image_uri("runtime-sandbox-appworld", arm64_compatible=False),
     sandbox_type=SandboxType.APPWORLD,
     runtime_config={"shm_size": "5.06gb"},
     security_level="medium",
@@ -258,7 +250,7 @@ DATASET_SUB_TYPE = os.environ.get("DATASET_SUB_TYPE", "multi_turn")
 
 
 @SandboxRegistry.register(
-    f"agentscope/runtime-sandbox-bfcl:{get_image_tag()}",
+    build_image_uri("runtime-sandbox-bfcl", arm64_compatible=False),
     sandbox_type=SandboxType.BFCL,
     runtime_config={"shm_size": "8.06gb"},
     security_level="medium",
@@ -309,7 +301,7 @@ class BFCLSandbox(TrainingSandbox):
 
 
 @SandboxRegistry.register(
-    f"agentscope/runtime-sandbox-webshop:{get_image_tag()}",
+    build_image_uri("runtime-sandbox-webshop", arm64_compatible=False),
     sandbox_type=SandboxType.WEBSHOP,
     runtime_config={"shm_size": "5.06gb"},
     security_level="medium",

agentscope_runtime/sandbox/build.py

@@ -11,6 +11,10 @@ import requests
 
 from .enums import SandboxType
 from .registry import SandboxRegistry
+from .utils import dynamic_import
+
+
+logger = logging.getLogger(__name__)
 
 
 def find_free_port(start_port, end_port):
@@ -18,7 +22,7 @@ def find_free_port(start_port, end_port):
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
             if sock.connect_ex(("localhost", port)) != 0:
                 return port
-    logging.error(
+    logger.error(
         f"No free ports available in the range {start_port}-{end_port}",
     )
     raise RuntimeError(
@@ -30,7 +34,7 @@ def check_health(url, secret_token, timeout=120, interval=5):
     headers = {"Authorization": f"Bearer {secret_token}"}
     spent_time = 0
     while spent_time < timeout:
-        logging.info(
+        logger.info(
             f"Attempting to connect to {url} (Elapsed time: {spent_time} "
             f"seconds)...",
         )
@@ -41,13 +45,13 @@ def check_health(url, secret_token, timeout=120, interval=5):
                 return True
         except requests.exceptions.RequestException:
             pass
-        logging.info(
+        logger.info(
             f"Health check failed for {url}. Retrying in {interval} "
             f"seconds...",
         )
         time.sleep(interval)
         spent_time += interval
-    logging.error(f"Health check failed for {url} after {timeout} seconds.")
+    logger.error(f"Health check failed for {url} after {timeout} seconds.")
     return False
 
 
@@ -57,10 +61,10 @@ def build_image(build_type, dockerfile_path=None):
             f"src/agentscope_runtime/sandbox/box/{build_type}/Dockerfile"
         )
 
-    logging.info(f"Building {build_type} with `{dockerfile_path}`...")
+    logger.info(f"Building {build_type} with `{dockerfile_path}`...")
 
     # Initialize and update Git submodule
-    logging.info("Initializing and updating Git submodule...")
+    logger.info("Initializing and updating Git submodule...")
     subprocess.run(
         ["git", "submodule", "update", "--init", "--recursive"],
         check=True,
@@ -69,7 +73,7 @@ def build_image(build_type, dockerfile_path=None):
     secret_token = "secret_token123"
     image_name = SandboxRegistry.get_image_by_type(build_type)
 
-    logging.info(f"Building Docker image {image_name}...")
+    logger.info(f"Building Docker image {image_name}...")
 
     # Check if image exists
     result = subprocess.run(
@@ -87,7 +91,7 @@ def build_image(build_type, dockerfile_path=None):
             f"you want to overwrite it? (y/N): ",
         )
         if choice.lower() != "y":
-            logging.info("Exiting without overwriting the existing image.")
+            logger.info("Exiting without overwriting the existing image.")
             return
 
     if not os.path.exists(dockerfile_path):
@@ -109,9 +113,9 @@ def build_image(build_type, dockerfile_path=None):
         ],
         check=False,
     )
-    logging.info(f"Docker image {image_name}dev built successfully.")
+    logger.info(f"Docker image {image_name}dev built successfully.")
 
-    logging.info(f"Start to build image {image_name}.")
+    logger.info(f"Start to build image {image_name}.")
 
     # Run the container with port mapping and environment variable
     free_port = find_free_port(8080, 8090)
@@ -131,7 +135,7 @@ def build_image(build_type, dockerfile_path=None):
         check=False,
     )
     container_id = result.stdout.strip()
-    logging.info(f"Running container {container_id} on port {free_port}")
+    logger.info(f"Running container {container_id} on port {free_port}")
 
     # Check health endpoints
     fastapi_health_url = f"http://localhost:{free_port}/fastapi/healthz"
@@ -149,18 +153,18 @@ def build_image(build_type, dockerfile_path=None):
         browser_healthy = True
 
     if browser_healthy and fastapi_healthy:
-        logging.info("Health checks passed.")
+        logger.info("Health checks passed.")
         subprocess.run(
             ["docker", "commit", container_id, f"{image_name}"],
             check=True,
         )
-        logging.info(
+        logger.info(
             f"Docker image {image_name} committed successfully.",
         )
         subprocess.run(["docker", "stop", container_id], check=True)
         subprocess.run(["docker", "rm", container_id], check=True)
     else:
-        logging.error("Health checks failed.")
+        logger.error("Health checks failed.")
         subprocess.run(["docker", "stop", container_id], check=True)
 
     choice = input(
@@ -171,9 +175,9 @@ def build_image(build_type, dockerfile_path=None):
             ["docker", "rmi", "-f", f"{image_name}dev"],
             check=True,
         )
-        logging.info(f"Dev image {image_name}dev deleted.")
+        logger.info(f"Dev image {image_name}dev deleted.")
     else:
-        logging.info(f"Dev image {image_name}dev retained.")
+        logger.info(f"Dev image {image_name}dev retained.")
 
 
 def main():
@@ -182,8 +186,8 @@ def main():
     )
     parser.add_argument(
         "build_type",
+        nargs="?",
         default="base",
-        choices=[x.value for x in SandboxType] + ["all"],
         help="Specify the build type to execute.",
     )
 
@@ -193,13 +197,29 @@ def main():
         help="Specify the path for the Dockerfile.",
     )
 
+    parser.add_argument(
+        "--extension",
+        action="append",
+        help="Path to a Python file or module name to load as an extension",
+    )
+
     args = parser.parse_args()
 
+    if args.extension:
+        for ext in args.extension:
+            logger.info(f"Loading extension: {ext}")
+            mod = dynamic_import(ext)
+            logger.info(f"Extension loaded: {mod.__name__}")
+
     if args.build_type == "all":
         # Only build the built-in images
         for build_type in [x.value for x in SandboxType.get_builtin_members()]:
             build_image(build_type)
     else:
+        assert args.build_type in [
+            x.value for x in SandboxType
+        ], f"Invalid build type: {args.build_type}"
+
         if args.build_type not in [
             x.value for x in SandboxType.get_builtin_members()
         ]:

agentscope_runtime/sandbox/client/http_client.py

@@ -69,7 +69,7 @@ class SandboxHttpClient:
     def __init__(
         self,
         model: Optional[ContainerModel] = None,
-        timeout: int = 30,
+        timeout: int = 60,
         enable_browser: bool = True,
         domain: str = "localhost",
     ) -> None:
@@ -95,7 +95,10 @@ class SandboxHttpClient:
         self.secret = model.runtime_token
 
         # Update headers with secret if provided
-        headers = {"Content-Type": "application/json"}
+        headers = {
+            "Content-Type": "application/json",
+            "x-agentrun-session-id": "s" + self.session_id,
+        }
         if self.secret:
             headers["Authorization"] = f"Bearer {self.secret}"
         self.session.headers.update(headers)
@@ -112,16 +115,20 @@ class SandboxHttpClient:
                 base_url=self.browser_url,
             )
 
+            headers = {"x-agentrun-session-id": "s" + self.session_id}
+
             # Create a new browser session if it doesn't exist
             try:
                 # Try to connet to existing session
                 self.steel_client.sessions.retrieve(
                     BROWSER_SESSION_ID,
+                    extra_headers=headers,
                 )
             except Exception:
                 # Session not found, create a new one
                 self.steel_client.sessions.create(
                     session_id=BROWSER_SESSION_ID,
+                    extra_headers=headers,
                 )
 
         return self
@@ -246,7 +253,10 @@ class SandboxHttpClient:
         """Run an IPython cell."""
         try:
             endpoint = f"{self.base_url}/tools/run_ipython_cell"
-            response = self.session.post(endpoint, json={"code": code})
+            response = self.session.post(
+                endpoint,
+                json={"code": code},
+            )
             response.raise_for_status()
             return response.json()
         except requests.exceptions.RequestException as e:
@@ -379,7 +389,11 @@ class SandboxHttpClient:
             endpoint = f"{self.base_url}/workspace/files"
             params = {"file_path": file_path}
             data = {"content": content}
-            response = self.session.post(endpoint, params=params, json=data)
+            response = self.session.post(
+                endpoint,
+                params=params,
+                json=data,
+            )
             response.raise_for_status()
             return response.json()
         except requests.exceptions.RequestException as e:
@@ -402,7 +416,10 @@ class SandboxHttpClient:
         try:
             endpoint = f"{self.base_url}/workspace/list-directories"
             params = {"directory": directory}
-            response = self.session.get(endpoint, params=params)
+            response = self.session.get(
+                endpoint,
+                params=params,
+            )
             response.raise_for_status()
             return response.json()
         except requests.exceptions.RequestException as e:
@@ -419,7 +436,10 @@ class SandboxHttpClient:
         try:
             endpoint = f"{self.base_url}/workspace/directories"
             params = {"directory_path": directory_path}
-            response = self.session.post(endpoint, params=params)
+            response = self.session.post(
+                endpoint,
+                params=params,
+            )
             response.raise_for_status()
             return response.json()
         except requests.exceptions.RequestException as e:
@@ -438,7 +458,10 @@ class SandboxHttpClient:
         try:
             endpoint = f"{self.base_url}/workspace/files"
             params = {"file_path": file_path}
-            response = self.session.delete(endpoint, params=params)
+            response = self.session.delete(
+                endpoint,
+                params=params,
+            )
             response.raise_for_status()
             return response.json()
         except requests.exceptions.RequestException as e:
@@ -461,7 +484,10 @@ class SandboxHttpClient:
         try:
             endpoint = f"{self.base_url}/workspace/directories"
             params = {"directory_path": directory_path, "recursive": recursive}
-            response = self.session.delete(endpoint, params=params)
+            response = self.session.delete(
+                endpoint,
+                params=params,
+            )
             response.raise_for_status()
             return response.json()
         except requests.exceptions.RequestException as e:
@@ -487,7 +513,10 @@ class SandboxHttpClient:
                 "source_path": source_path,
                 "destination_path": destination_path,
             }
-            response = self.session.put(endpoint, params=params)
+            response = self.session.put(
+                endpoint,
+                params=params,
+            )
             response.raise_for_status()
             return response.json()
         except requests.exceptions.RequestException as e:
@@ -514,7 +543,10 @@ class SandboxHttpClient:
                 "source_path": source_path,
                 "destination_path": destination_path,
             }
-            response = self.session.post(endpoint, params=params)
+            response = self.session.post(
+                endpoint,
+                params=params,
+            )
             response.raise_for_status()
             return response.json()
         except requests.exceptions.RequestException as e:

agentscope_runtime/sandbox/client/training_client.py

@@ -212,7 +212,6 @@ class TrainingSandboxClient:
     def list_tools(self, **kwargs):
         """list tools"""
         if "instance_id" in kwargs:
-            # 只传get_tools_info的方法参数
             return self.get_tools_info(
                 instance_id=kwargs.get("instance_id"),
                 messages=kwargs.get("messages", {}),

agentscope_runtime/sandbox/constant.py

@@ -1,5 +1,31 @@
 # -*- coding: utf-8 -*-
 import os
+import sys
+import logging
 
+logger = logging.getLogger(__name__)
+
+# Image Registry: default "", means using Docker Hub.
+# For AgentScope official registry, use
+# "agentscope-registry.ap-southeast-1.cr.aliyuncs.com"
+REGISTRY = os.getenv("RUNTIME_SANDBOX_REGISTRY", "")
+if REGISTRY == "":
+    agentscope_acr = "agentscope-registry.ap-southeast-1.cr.aliyuncs.com"
+    if sys.platform.startswith("win"):
+        cmd = f"set RUNTIME_SANDBOX_REGISTRY={agentscope_acr}"
+    else:
+        cmd = f"export RUNTIME_SANDBOX_REGISTRY={agentscope_acr}"
+    logger.warning(
+        "Using Docker Hub as image registry. If pulling is slow or fails, "
+        f"you can switch to the AgentScope official registry by running:\n  "
+        f"{cmd}\n",
+    )
+
+# Image Namespace
+IMAGE_NAMESPACE = os.getenv("RUNTIME_SANDBOX_IMAGE_NAMESPACE", "agentscope")
+
+# Image Tag
 IMAGE_TAG = os.getenv("RUNTIME_SANDBOX_IMAGE_TAG", "latest")
+
+# Browser Session ID
 BROWSER_SESSION_ID = "123e4567-e89b-12d3-a456-426614174000"

agentscope_runtime/sandbox/custom/custom_sandbox.py

@@ -3,17 +3,17 @@ import os
 
 from typing import Optional
 
-from ..constant import IMAGE_TAG
+from ..utils import build_image_uri
 from ..registry import SandboxRegistry
 from ..enums import SandboxType
 from ..box.sandbox import Sandbox
 
-SANDBOXTYPE = "custom_sandbox"
+SANDBOX_TYPE = "custom_sandbox"
 
 
 @SandboxRegistry.register(
-    f"agentscope/runtime-sandbox-{SANDBOXTYPE}:{IMAGE_TAG}",
-    sandbox_type=SANDBOXTYPE,
+    build_image_uri(f"runtime-sandbox-{SANDBOX_TYPE}"),
+    sandbox_type=SANDBOX_TYPE,
     security_level="medium",
     timeout=60,
     description="my sandbox",
@@ -35,5 +35,5 @@ class CustomSandbox(Sandbox):
             timeout,
             base_url,
             bearer_token,
-            SandboxType(SANDBOXTYPE),
+            SandboxType(SANDBOX_TYPE),
         )

agentscope_runtime/sandbox/custom/example.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 from typing import Optional
 
-from ..constant import IMAGE_TAG
+from ..utils import build_image_uri
 from ..registry import SandboxRegistry
 from ..enums import SandboxType
 from ..box.sandbox import Sandbox
@@ -10,7 +10,7 @@ SANDBOX_TYPE = "example"
 
 
 @SandboxRegistry.register(
-    f"agentscope/runtime-sandbox-{SANDBOX_TYPE}:{IMAGE_TAG}",
+    build_image_uri(f"runtime-sandbox-{SANDBOX_TYPE}"),
     sandbox_type=SANDBOX_TYPE,
     security_level="medium",
     timeout=60,

agentscope_runtime/sandbox/manager/collections/in_memory_mapping.py

@@ -1,4 +1,6 @@
 # -*- coding: utf-8 -*-
+from typing import Any
+
 from .base_mapping import Mapping
 
 
@@ -6,10 +8,10 @@ class InMemoryMapping(Mapping):
     def __init__(self):
         self.store = {}
 
-    def set(self, key: str, value: dict):
+    def set(self, key: str, value: Any):
         self.store[key] = value
 
-    def get(self, key: str) -> dict:
+    def get(self, key: str) -> Any:
         return self.store.get(key)
 
     def delete(self, key: str):

agentscope_runtime/sandbox/manager/collections/redis_mapping.py

@@ -1,26 +1,42 @@
 # -*- coding: utf-8 -*-
 import json
 
+from typing import Any
+
 from .base_mapping import Mapping
 
 
 class RedisMapping(Mapping):
-    def __init__(self, redis_client):
+    def __init__(self, redis_client, prefix: str = ""):
         self.client = redis_client
+        self.prefix = prefix.rstrip(":") + ":" if prefix else ""
+
+    def _get_full_key(self, key: str) -> str:
+        return f"{self.prefix}{key}"
+
+    def _strip_prefix(self, full_key: str) -> str:
+        if self.prefix and full_key.startswith(self.prefix):
+            return full_key[len(self.prefix) :]
+        return full_key
 
-    def set(self, key: str, value: dict):
-        self.client.set(key, json.dumps(value))
+    def set(self, key: str, value: Any):
+        self.client.set(self._get_full_key(key), json.dumps(value))
 
-    def get(self, key: str) -> dict:
-        value = self.client.get(key)
+    def get(self, key: str) -> Any:
+        value = self.client.get(self._get_full_key(key))
         return json.loads(value) if value else None
 
     def delete(self, key: str):
-        self.client.delete(key)
+        self.client.delete(self._get_full_key(key))
 
-    def scan(self, prefix: str):
+    def scan(self, prefix: str = ""):
+        search_pattern = f"{self._get_full_key(prefix)}*"
         cursor = 0
         while cursor != 0:
-            cursor, keys = self.client.scan(cursor=cursor, match=f"{prefix}*")
+            cursor, keys = self.client.scan(
+                cursor=cursor,
+                match=search_pattern,
+            )
             for key in keys:
-                yield key.decode("utf-8")
+                decoded_key = key.decode("utf-8")
+                yield self._strip_prefix(decoded_key)

agentscope_runtime/sandbox/manager/container_clients/__init__.py

@@ -1,10 +0,0 @@
-# -*- coding: utf-8 -*-
-from .base_client import BaseClient
-from .docker_client import DockerClient
-from .kubernetes_client import KubernetesClient
-
-__all__ = [
-    "BaseClient",
-    "DockerClient",
-    "KubernetesClient",
-]