File size: 3,528 Bytes
44e7e06
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
from typing import Dict, List, Any, Optional

from aworld.sandbox.base import Sandbox
from aworld.sandbox.common import BaseSandbox
from aworld.sandbox.models import SandboxEnvType
from aworld.sandbox.implementations import LocalSandbox, KubernetesSandbox, SuperSandbox


# For backward compatibility, use LocalSandbox as the default Sandbox implementation
DefaultSandbox = LocalSandbox


# Override Sandbox class constructor to create the appropriate sandbox based on env_type
def create_sandbox(
    env_type: Optional[int] = None,
    sandbox_id: Optional[str] = None,
    metadata: Optional[Dict[str, str]] = None,
    timeout: Optional[int] = None,
    mcp_servers: Optional[List[str]] = None,
    mcp_config: Optional[Any] = None,
    **kwargs
) -> Sandbox:
    """
    Factory function to create a sandbox instance based on the environment type.
    
    Args:
        env_type: The environment type. Defaults to LOCAL if None.
        sandbox_id: Unique identifier for the sandbox. If None, one will be generated.
        metadata: Additional metadata for the sandbox.
        timeout: Timeout for sandbox operations.
        mcp_servers: List of MCP servers to use.
        mcp_config: Configuration for MCP servers.
        **kwargs: Additional parameters for specific sandbox types.
        
    Returns:
        Sandbox: An instance of a sandbox implementation.
        
    Raises:
        ValueError: If an invalid environment type is provided.
    """
    env_type = env_type or SandboxEnvType.LOCAL
    
    if env_type == SandboxEnvType.LOCAL:
        return LocalSandbox(
            sandbox_id=sandbox_id,
            metadata=metadata,
            timeout=timeout,
            mcp_servers=mcp_servers,
            mcp_config=mcp_config,
            **kwargs
        )
    elif env_type == SandboxEnvType.K8S:
        return KubernetesSandbox(
            sandbox_id=sandbox_id,
            metadata=metadata,
            timeout=timeout,
            mcp_servers=mcp_servers,
            mcp_config=mcp_config,
            **kwargs
        )
    elif env_type == SandboxEnvType.SUPERCOMPUTER:
        return SuperSandbox(
            sandbox_id=sandbox_id,
            metadata=metadata,
            timeout=timeout,
            mcp_servers=mcp_servers,
            mcp_config=mcp_config,
            **kwargs
        )
    else:
        raise ValueError(f"Invalid environment type: {env_type}")


# Monkey patch the Sandbox class to make direct instantiation work
old_init = Sandbox.__init__

def _sandbox_init(self, *args, **kwargs):
    if type(self) is Sandbox:
        # This should never be called directly, as __new__ will return a different type
        pass
    else:
        # Pass through to the original __init__ for actual implementations
        old_init(self, *args, **kwargs)

# Store the original __new__ method
original_new = object.__new__

# Create a new __new__ method that intercepts Sandbox instantiation
def _sandbox_new(cls, *args, **kwargs):
    if cls is Sandbox:
        # If trying to instantiate Sandbox directly, use our factory instead
        return create_sandbox(**kwargs)
    else:
        # For subclasses, use the original __new__
        return original_new(cls)

# Apply the monkey patches
Sandbox.__init__ = _sandbox_init
Sandbox.__new__ = _sandbox_new


# Expose key classes and functions
__all__ = [
    'Sandbox',
    'BaseSandbox',
    'LocalSandbox',
    'KubernetesSandbox',
    'SuperSandbox',
    'DefaultSandbox',
    'SandboxEnvType',
    'create_sandbox'
]