diff --git a/api/core/virtual_environment/__base/virtual_environment.py b/api/core/virtual_environment/__base/virtual_environment.py index 439b0aa085..f006e296d6 100644 --- a/api/core/virtual_environment/__base/virtual_environment.py +++ b/api/core/virtual_environment/__base/virtual_environment.py @@ -18,10 +18,10 @@ class VirtualEnvironment(ABC): """ self.options = options - self.metadata = self.construct_environment(options, environments or {}) + self.metadata = self._construct_environment(options, environments or {}) @abstractmethod - def construct_environment(self, options: Mapping[str, Any], environments: Mapping[str, str]) -> Metadata: + def _construct_environment(self, options: Mapping[str, Any], environments: Mapping[str, str]) -> Metadata: """ Construct the unique identifier for the virtual environment. diff --git a/api/core/virtual_environment/providers/docker_daemon_sandbox.py b/api/core/virtual_environment/providers/docker_daemon_sandbox.py index bb9b626756..3e52c9feee 100644 --- a/api/core/virtual_environment/providers/docker_daemon_sandbox.py +++ b/api/core/virtual_environment/providers/docker_daemon_sandbox.py @@ -77,7 +77,7 @@ class DockerDaemonEnvironment(VirtualEnvironment): DOCKER_IMAGE = "docker_image" DOCKER_COMMAND = "docker_command" - def construct_environment(self, options: Mapping[str, Any], environments: Mapping[str, str]) -> Metadata: + def _construct_environment(self, options: Mapping[str, Any], environments: Mapping[str, str]) -> Metadata: """ Construct the Docker daemon virtual environment. """ diff --git a/api/core/virtual_environment/providers/e2b_sandbox.py b/api/core/virtual_environment/providers/e2b_sandbox.py index 86fa58fb68..3d777fb7be 100644 --- a/api/core/virtual_environment/providers/e2b_sandbox.py +++ b/api/core/virtual_environment/providers/e2b_sandbox.py @@ -85,7 +85,7 @@ class E2BEnvironment(VirtualEnvironment): class StoreKey(StrEnum): SANDBOX = "sandbox" - def construct_environment(self, options: Mapping[str, Any], environments: Mapping[str, str]) -> Metadata: + def _construct_environment(self, options: Mapping[str, Any], environments: Mapping[str, str]) -> Metadata: """ Construct a new E2B virtual environment. """ diff --git a/api/core/virtual_environment/providers/local_without_isolation.py b/api/core/virtual_environment/providers/local_without_isolation.py index c8641e64c7..29ec8c0569 100644 --- a/api/core/virtual_environment/providers/local_without_isolation.py +++ b/api/core/virtual_environment/providers/local_without_isolation.py @@ -65,7 +65,7 @@ class LocalVirtualEnvironment(VirtualEnvironment): NEVER USE IT IN PRODUCTION ENVIRONMENTS. """ - def construct_environment(self, options: Mapping[str, Any], environments: Mapping[str, str]) -> Metadata: + def _construct_environment(self, options: Mapping[str, Any], environments: Mapping[str, str]) -> Metadata: """ Construct the local virtual environment. diff --git a/api/core/virtual_environment/session/__init__.py b/api/core/virtual_environment/session/__init__.py deleted file mode 100644 index a4c24ce169..0000000000 --- a/api/core/virtual_environment/session/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .sandbox_session import SandboxProvider, SandboxSession, SandboxSessionManager - -__all__ = ["SandboxProvider", "SandboxSession", "SandboxSessionManager"] diff --git a/api/core/virtual_environment/session/sandbox_session.py b/api/core/virtual_environment/session/sandbox_session.py deleted file mode 100644 index d2bf64d7eb..0000000000 --- a/api/core/virtual_environment/session/sandbox_session.py +++ /dev/null @@ -1,47 +0,0 @@ -from enum import StrEnum -from typing import Any - -from pydantic import Field - -from core.session import BaseSession, SessionManager -from core.virtual_environment.__base.entities import Arch - - -class SandboxProvider(StrEnum): - E2B = "e2b" - DOCKER = "docker" - LOCAL = "local" - - -class SandboxSession(BaseSession): - provider: SandboxProvider - sandbox_id: str - arch: Arch - connection_config: dict[str, Any] = Field(default_factory=dict) - - -class SandboxSessionManager(SessionManager[SandboxSession]): - def __init__(self, ttl: int | None = None): - super().__init__(key_prefix="sandbox_session", session_class=SandboxSession, ttl=ttl) - - def create( - self, - tenant_id: str, - user_id: str, - provider: SandboxProvider, - sandbox_id: str, - arch: Arch, - connection_config: dict[str, Any] | None = None, - context: dict[str, Any] | None = None, - ) -> SandboxSession: - session = SandboxSession( - tenant_id=tenant_id, - user_id=user_id, - provider=provider, - sandbox_id=sandbox_id, - arch=arch, - connection_config=connection_config or {}, - context=context or {}, - ) - self.save(session) - return session