Simplify connection config handling in Python SDK (#849)

Several improvements and refactors to the Python SDK's sandbox
connection and configuration logic.

The main focus is on making header management more flexible, improving
API parameter handling, and simplifying class structures for better
maintainability and usability.

Also unifies doc strings

Author: jakubno

.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