Simplify connection config handling in Python SDK

.changeset/strange-hornets-judge.md

@@ -0,0 +1,5 @@
+---
+'@e2b/python-sdk': minor
+---
+
+Refactor connection config and parameters for all calls to API

packages/python-sdk/e2b/connection_config.py

@@ -1,7 +1,8 @@
 import os
 
-from typing import Literal, Optional, Dict
+from typing import Literal, Optional, Dict, TypedDict
 from httpx._types import ProxyTypes
+from typing_extensions import Unpack
 
 from e2b.api.metadata import package_version
 
@@ -11,6 +12,32 @@
 KEEPALIVE_PING_HEADER = "Keepalive-Ping-Interval"
 
 
+class ApiParams(TypedDict, total=False):
+    """
+    Parameters for a request.
+
+    In the case of a sandbox, it applies to all **requests made to the returned sandbox**.
+    """
+
+    request_timeout: Optional[float]
+    """Timeout for the request in **seconds**, defaults to 60 seconds."""
+
+    headers: Optional[Dict[str, str]]
+    """Additional headers to send with the request."""
+
+    api_key: Optional[str]
+    """E2B API Key to use for authentication, defaults to `E2B_API_KEY` environment variable."""
+
+    domain: Optional[str]
+    """E2B domain to use for authentication, defaults to `E2B_DOMAIN` environment variable."""
+
+    debug: Optional[bool]
+    """Whether to use debug mode, defaults to `E2B_DEBUG` environment variable."""
+
+    proxy: Optional[ProxyTypes]
+    """Proxy to use for the request. In case of a sandbox it applies to all **requests made to the returned sandbox**."""
+
+
 class ConnectionConfig:
     """
     Configuration for the connection to the API.
@@ -40,6 +67,7 @@ def __init__(
         access_token: Optional[str] = None,
         request_timeout: Optional[float] = None,
         headers: Optional[Dict[str, str]] = None,
+        extra_sandbox_headers: Optional[Dict[str, str]] = None,
         proxy: Optional[ProxyTypes] = None,
     ):
         self.domain = domain or ConnectionConfig._domain()
@@ -48,6 +76,7 @@ def __init__(
         self.access_token = access_token or ConnectionConfig._access_token()
         self.headers = headers or {}
         self.headers["User-Agent"] = f"e2b-python-sdk/{package_version}"
+        self.__extra_sandbox_headers = extra_sandbox_headers or {}
 
         self.proxy = proxy
 
@@ -82,6 +111,53 @@ def _get_request_timeout(
     def get_request_timeout(self, request_timeout: Optional[float] = None):
         return self._get_request_timeout(self.request_timeout, request_timeout)
 
+    def get_api_params(
+        self,
+        **opts: Unpack[ApiParams],
+    ) -> dict:
+        """
+        Get the parameters for the API call.
+
+        This is used to avoid passing the following attributes to the API call:
+        - access_token
+        - api_url
+
+        It also returns a copy, so the original object is not modified.
+
+        :return: Dictionary of parameters for the API call
+        """
+        headers = opts.get("headers")
+        request_timeout = opts.get("request_timeout")
+        api_key = opts.get("api_key")
+        domain = opts.get("domain")
+        debug = opts.get("debug")
+        proxy = opts.get("proxy")
+
+        req_headers = self.headers.copy()
+        if headers is not None:
+            req_headers.update(headers)
+
+        return dict(
+            ApiParams(
+                api_key=api_key if api_key is not None else self.api_key,
+                domain=domain if domain is not None else self.domain,
+                debug=debug if debug is not None else self.debug,
+                request_timeout=self.get_request_timeout(request_timeout),
+                headers=req_headers,
+                proxy=proxy if proxy is not None else self.proxy,
+            )
+        )
+
+    @property
+    def sandbox_headers(self):
+        """
+        # We need this separate as we use the same header for E2B access token to API and envd access token to sandbox.
+        """
+        return {
+            **self.headers,
+            **self.__extra_sandbox_headers,
+        }
+
 
 Username = Literal["root", "user"]
 """

packages/python-sdk/e2b/sandbox/main.py

@@ -1,6 +1,5 @@
 import urllib.parse
 
-from abc import ABC, abstractmethod
 from typing import Optional
 
 from e2b.sandbox.signature import get_signature
@@ -9,7 +8,7 @@
 from httpx import Limits
 
 
-class SandboxSetup(ABC):
+class SandboxBase:
     _limits = Limits(
         max_keepalive_connections=40,
         max_connections=40,
@@ -21,30 +20,48 @@ class SandboxSetup(ABC):
     default_sandbox_timeout = 300
     default_template = "base"
 
+    def __init__(
+        self,
+        sandbox_id: str,
+        envd_version: Optional[str],
+        envd_access_token: Optional[str],
+        sandbox_domain: Optional[str],
+        connection_config: ConnectionConfig,
+    ):
+        self.__connection_config = connection_config
+        self.__sandbox_id = sandbox_id
+        self.__sandbox_domain = sandbox_domain or self.connection_config.domain
+        self.__envd_version = envd_version
+        self.__envd_access_token = envd_access_token
+        self.__envd_api_url = f"{'http' if self.connection_config.debug else 'https'}://{self.get_host(self.envd_port)}"
+
+    @property
+    def _envd_access_token(self) -> Optional[str]:
+        """Private property to access the envd token"""
+        return self.__envd_access_token
+
     @property
-    @abstractmethod
     def connection_config(self) -> ConnectionConfig:
-        ...
+        return self.__connection_config
 
     @property
-    @abstractmethod
-    def _envd_access_token(self) -> Optional[str]:
-        ...
+    def _envd_version(self) -> Optional[str]:
+        return self.__envd_version
 
     @property
-    @abstractmethod
-    def envd_api_url(self) -> str:
-        ...
+    def sandbox_domain(self) -> Optional[str]:
+        return self.__sandbox_domain
 
     @property
-    @abstractmethod
-    def sandbox_id(self) -> str:
-        ...
+    def envd_api_url(self) -> str:
+        return self.__envd_api_url
 
     @property
-    @abstractmethod
-    def sandbox_domain(self) -> str:
-        ...
+    def sandbox_id(self) -> str:
+        """
+        Unique identifier of the sandbox.
+        """
+        return self.__sandbox_id
 
     def _file_url(
         self,

packages/python-sdk/e2b/sandbox/sandbox_api.py

@@ -1,10 +1,7 @@
-from abc import ABC
 from dataclasses import dataclass
 from typing import Optional, Dict
 from datetime import datetime
 
-from httpx import Limits
-
 from e2b.api.client.models import SandboxState
 
 
@@ -81,11 +78,3 @@ class SandboxMetrics:
     """Memory used in bytes."""
     timestamp: datetime
     """Timestamp of the metric entry."""
-
-
-class SandboxApiBase(ABC):
-    _limits = Limits(
-        max_keepalive_connections=10,
-        max_connections=20,
-        keepalive_expiry=20,
-    )

packages/python-sdk/e2b/sandbox_async/commands/command.py

@@ -35,7 +35,7 @@ def __init__(
             # compressor=e2b_connect.GzipCompressor,
             async_pool=pool,
             json=True,
-            headers=connection_config.headers,
+            headers=connection_config.sandbox_headers,
         )
 
     async def list(

packages/python-sdk/e2b/sandbox_async/commands/pty.py

@@ -39,6 +39,7 @@ def __init__(
             # compressor=e2b_connect.GzipCompressor,
             async_pool=pool,
             json=True,
+            headers=connection_config.sandbox_headers,
         )
 
     async def kill(

packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py

@@ -51,7 +51,7 @@ def __init__(
             # compressor=e2b_connect.GzipCompressor,
             async_pool=pool,
             json=True,
-            headers=connection_config.headers,
+            headers=connection_config.sandbox_headers,
         )
 
     @overload