aip-agents-binary 0.6.5__py3-none-macosx_13_0_arm64.whl → 0.6.6__py3-none-macosx_13_0_arm64.whl

aip_agents/ptc/executor.py

@@ -1,7 +1,6 @@
-"""PTC Executor implementations (MCP-only).
+"""PTC Executor implementations.
 
-This module provides the sandboxed executor for Programmatic Tool Calling
-with MCP tools.
+This module provides the sandboxed executor for Programmatic Tool Calling.
 
 Authors:
     Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
@@ -9,9 +8,12 @@ Authors:
 
 from __future__ import annotations
 
+import asyncio
 from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
 
 from aip_agents.mcp.client.base_mcp_client import BaseMCPClient
+from aip_agents.ptc.custom_tools import PTCCustomToolConfig
 from aip_agents.ptc.exceptions import PTCToolError
 from aip_agents.ptc.prompt_builder import PromptConfig
 from aip_agents.sandbox.defaults import DEFAULT_PTC_PACKAGES, DEFAULT_PTC_TEMPLATE
@@ -30,18 +32,22 @@ except ImportError:
 
 logger = get_logger(__name__)
 
+if TYPE_CHECKING:
+    from aip_agents.ptc.payload import SandboxPayload
+
 
 @dataclass
 class PTCSandboxConfig:
-    """Configuration for PTC sandbox executor (MCP-only).
+    """Configuration for PTC sandbox executor.
 
     Attributes:
         enabled: Whether PTC is enabled. When False, PTC is disabled.
         default_tool_timeout: Default timeout per tool call in seconds.
         sandbox_template: Optional E2B sandbox template ID.
         sandbox_timeout: Sandbox execution timeout in seconds (hard cap/TTL).
-        ptc_packages: List of packages to install in sandbox. None or empty list skips install.
+        ptc_packages: List of packages to install in sandbox. Defaults to DEFAULT_PTC_PACKAGES.
         prompt: Prompt configuration for PTC usage guidance.
+        custom_tools: Configuration for custom LangChain tools in sandbox.
     """
 
     enabled: bool = False
@@ -50,31 +56,46 @@ class PTCSandboxConfig:
     sandbox_timeout: float = 300.0
     ptc_packages: list[str] | None = field(default_factory=lambda: list(DEFAULT_PTC_PACKAGES))
     prompt: PromptConfig = field(default_factory=PromptConfig)
+    custom_tools: PTCCustomToolConfig = field(default_factory=PTCCustomToolConfig)
 
 
 class PTCSandboxExecutor:
-    r"""Executes PTC code inside an E2B sandbox (MCP-only).
+    r"""Executes PTC code inside an E2B sandbox.
 
     This executor is used for LLM-generated code that requires sandboxing.
     It builds a sandbox payload (MCP server config + generated tool modules)
     and executes the code using the E2B runtime.
 
+    Static bundle caching:
+        The executor tracks whether the static bundle (wrappers, registry, sources)
+        has been uploaded. On the first run, it uploads the full payload. On subsequent
+        runs, it only uploads per-run files (e.g., tools/custom_defaults.json) to
+        reduce upload overhead.
+
+        If the sandbox is destroyed/recreated, call reset_static_bundle() to force
+        a full re-upload on the next execution.
+
+    Thread-safety:
+        The static bundle upload is guarded by an asyncio.Lock. Concurrent calls will
+        serialize while the initial upload completes. After the bundle is cached,
+        executions proceed without locking.
+
     Example:
         runtime = E2BSandboxRuntime()
         executor = PTCSandboxExecutor(mcp_client, runtime)
-        result = await executor.execute_code("from tools.yfinance import get_stock\nprint(get_stock('AAPL'))")
+        result = await executor.execute_code("from tools.yfinance import get_stock\\nprint(get_stock('AAPL'))")
     """
 
     def __init__(
         self,
-        mcp_client: BaseMCPClient,
+        mcp_client: BaseMCPClient | None,
         runtime: E2BSandboxRuntime,
         config: PTCSandboxConfig | None = None,
     ) -> None:
         """Initialize PTCSandboxExecutor.
 
         Args:
-            mcp_client: The MCP client with configured servers.
+            mcp_client: The MCP client with configured servers. Can be None for custom-only configs.
             runtime: The E2B sandbox runtime instance.
             config: Optional sandbox executor configuration.
 
@@ -90,21 +111,127 @@ class PTCSandboxExecutor:
         self._mcp_client = mcp_client
         self._runtime = runtime
         self._config = config or PTCSandboxConfig()
+        self._static_bundle_uploaded = False
+        self._bundle_lock = asyncio.Lock()
+
+    def _reset_bundle_state_if_inactive(self) -> None:
+        """Reset cached bundle state when the runtime is inactive."""
+        if self._runtime.is_active:
+            return
+        if not self._static_bundle_uploaded:
+            return
+
+        logger.info("Runtime is inactive, resetting static bundle state")
+        self._static_bundle_uploaded = False
+
+    def _should_include_packages_path(self) -> bool:
+        """Check if the packages path should be added to sys.path."""
+        custom_tools = self._config.custom_tools
+        if not custom_tools.enabled or not custom_tools.tools:
+            return False
+
+        return any(tool.get("package_path") or tool.get("kind") == "file" for tool in custom_tools.tools)
+
+    async def _build_payload(self, tool_configs: dict[str, dict] | None) -> SandboxPayload:
+        """Build the sandbox payload for execution."""
+        logger.info("Building sandbox payload")
+        return await build_sandbox_payload(
+            self._mcp_client,
+            self._config.default_tool_timeout,
+            custom_tools_config=self._config.custom_tools,
+            tool_configs=tool_configs,
+        )
+
+    def _wrap_code(self, code: str) -> str:
+        """Wrap code with required imports and setup."""
+        logger.info("Wrapping PTC code")
+        return wrap_ptc_code(code, include_packages_path=self._should_include_packages_path())
+
+    async def _execute_payload(
+        self,
+        payload: SandboxPayload,
+        wrapped_code: str,
+        *,
+        upload_static_bundle: bool,
+    ) -> SandboxExecutionResult:
+        """Execute the wrapped code with the provided payload."""
+        if upload_static_bundle:
+            logger.info("Uploading static bundle and per-run files")
+            files_to_upload = {**payload.files, **payload.per_run_files}
+        else:
+            logger.debug("Static bundle already uploaded, uploading only per-run files")
+            files_to_upload = payload.per_run_files
+
+        logger.info(f"Executing code in sandbox (timeout: {self._config.sandbox_timeout}s)")
+        return await self._runtime.execute(
+            code=wrapped_code,
+            timeout=self._config.sandbox_timeout,
+            files=files_to_upload if files_to_upload else None,
+            env=payload.env,
+            template=self._config.sandbox_template,
+        )
+
+    async def _execute_with_bundle(
+        self,
+        payload: SandboxPayload,
+        wrapped_code: str,
+    ) -> SandboxExecutionResult:
+        """Execute code using the cached static bundle when possible."""
+        if self._static_bundle_uploaded:
+            return await self._execute_payload(
+                payload,
+                wrapped_code,
+                upload_static_bundle=False,
+            )
+
+        async with self._bundle_lock:
+            if self._static_bundle_uploaded:
+                return await self._execute_payload(
+                    payload,
+                    wrapped_code,
+                    upload_static_bundle=False,
+                )
+
+            result = await self._execute_payload(
+                payload,
+                wrapped_code,
+                upload_static_bundle=True,
+            )
+            self._update_static_bundle_cache(result)
+            return result
+
+    def _update_static_bundle_cache(self, result: SandboxExecutionResult) -> None:
+        """Update cached bundle state after a successful upload."""
+        if result.exit_code != 0:
+            return
+
+        self._static_bundle_uploaded = True
+        logger.debug("Static bundle successfully uploaded and cached")
+
+    def _log_execution_result(self, result: SandboxExecutionResult) -> None:
+        """Log execution results based on exit status."""
+        if result.exit_code == 0:
+            logger.info("Sandbox execution completed successfully")
+        else:
+            logger.warning(f"Sandbox execution failed with exit code {result.exit_code}")
 
     async def execute_code(
         self,
         code: str,
+        tool_configs: dict[str, dict] | None = None,
     ) -> SandboxExecutionResult:
         """Execute code inside the sandbox with MCP access.
 
         This method:
-        1. Builds the sandbox payload (MCP config + generated tool modules)
+        1. Builds the sandbox payload (MCP config + generated tool modules + custom tools)
         2. Wraps the user code with necessary imports and setup
         3. Executes the code in the E2B sandbox
         4. Returns the execution result (stdout/stderr/exit_code)
 
         Args:
             code: Python code to execute in the sandbox.
+            tool_configs: Optional per-tool config values for custom LangChain tools.
+                These are merged with agent defaults and written to custom_defaults.json.
 
         Returns:
             SandboxExecutionResult with stdout, stderr, and exit_code.
@@ -113,31 +240,22 @@ class PTCSandboxExecutor:
             PTCToolError: If sandbox execution fails.
         """
         try:
-            logger.info("Building sandbox payload")
-            payload = await build_sandbox_payload(
-                self._mcp_client,
-                self._config.default_tool_timeout,
-            )
-
-            logger.info("Wrapping PTC code")
-            wrapped_code = wrap_ptc_code(code)
-
-            logger.info(f"Executing code in sandbox (timeout: {self._config.sandbox_timeout}s)")
-            result = await self._runtime.execute(
-                code=wrapped_code,
-                timeout=self._config.sandbox_timeout,
-                files=payload.files if payload.files else None,
-                env=payload.env,
-                template=self._config.sandbox_template,
-            )
-
-            if result.exit_code == 0:
-                logger.info("Sandbox execution completed successfully")
-            else:
-                logger.warning(f"Sandbox execution failed with exit code {result.exit_code}")
-
+            self._reset_bundle_state_if_inactive()
+            payload = await self._build_payload(tool_configs)
+            wrapped_code = self._wrap_code(code)
+            result = await self._execute_with_bundle(payload, wrapped_code)
+            self._log_execution_result(result)
             return result
 
         except Exception as exc:
             logger.error(f"Sandbox execution failed: {exc}")
             raise PTCToolError(f"Sandbox execution failed: {exc}") from exc
+
+    def reset_static_bundle(self) -> None:
+        """Reset the static bundle upload state.
+
+        Call this method when the sandbox is destroyed/recreated to force
+        a full re-upload of the static bundle on the next execution.
+        """
+        self._static_bundle_uploaded = False
+        logger.debug("Static bundle state reset, next execution will re-upload full payload")

aip_agents/ptc/executor.pyi

@@ -1,6 +1,8 @@
 from _typeshed import Incomplete
 from aip_agents.mcp.client.base_mcp_client import BaseMCPClient as BaseMCPClient
+from aip_agents.ptc.custom_tools import PTCCustomToolConfig as PTCCustomToolConfig
 from aip_agents.ptc.exceptions import PTCToolError as PTCToolError
+from aip_agents.ptc.payload import SandboxPayload as SandboxPayload
 from aip_agents.ptc.prompt_builder import PromptConfig as PromptConfig
 from aip_agents.ptc.sandbox_bridge import build_sandbox_payload as build_sandbox_payload, wrap_ptc_code as wrap_ptc_code
 from aip_agents.sandbox.defaults import DEFAULT_PTC_PACKAGES as DEFAULT_PTC_PACKAGES, DEFAULT_PTC_TEMPLATE as DEFAULT_PTC_TEMPLATE
@@ -13,15 +15,16 @@ logger: Incomplete
 
 @dataclass
 class PTCSandboxConfig:
-    """Configuration for PTC sandbox executor (MCP-only).
+    """Configuration for PTC sandbox executor.
 
     Attributes:
         enabled: Whether PTC is enabled. When False, PTC is disabled.
         default_tool_timeout: Default timeout per tool call in seconds.
         sandbox_template: Optional E2B sandbox template ID.
         sandbox_timeout: Sandbox execution timeout in seconds (hard cap/TTL).
-        ptc_packages: List of packages to install in sandbox. None or empty list skips install.
+        ptc_packages: List of packages to install in sandbox. Defaults to DEFAULT_PTC_PACKAGES.
         prompt: Prompt configuration for PTC usage guidance.
+        custom_tools: Configuration for custom LangChain tools in sandbox.
     """
     enabled: bool = ...
     default_tool_timeout: float = ...
@@ -29,41 +32,58 @@ class PTCSandboxConfig:
     sandbox_timeout: float = ...
     ptc_packages: list[str] | None = field(default_factory=Incomplete)
     prompt: PromptConfig = field(default_factory=PromptConfig)
+    custom_tools: PTCCustomToolConfig = field(default_factory=PTCCustomToolConfig)
 
 class PTCSandboxExecutor:
-    '''Executes PTC code inside an E2B sandbox (MCP-only).
+    '''Executes PTC code inside an E2B sandbox.
 
     This executor is used for LLM-generated code that requires sandboxing.
     It builds a sandbox payload (MCP server config + generated tool modules)
     and executes the code using the E2B runtime.
 
+    Static bundle caching:
+        The executor tracks whether the static bundle (wrappers, registry, sources)
+        has been uploaded. On the first run, it uploads the full payload. On subsequent
+        runs, it only uploads per-run files (e.g., tools/custom_defaults.json) to
+        reduce upload overhead.
+
+        If the sandbox is destroyed/recreated, call reset_static_bundle() to force
+        a full re-upload on the next execution.
+
+    Thread-safety:
+        The static bundle upload is guarded by an asyncio.Lock. Concurrent calls will
+        serialize while the initial upload completes. After the bundle is cached,
+        executions proceed without locking.
+
     Example:
         runtime = E2BSandboxRuntime()
         executor = PTCSandboxExecutor(mcp_client, runtime)
-        result = await executor.execute_code("from tools.yfinance import get_stock\\nprint(get_stock(\'AAPL\'))")
+        result = await executor.execute_code("from tools.yfinance import get_stock\\\\nprint(get_stock(\'AAPL\'))")
     '''
-    def __init__(self, mcp_client: BaseMCPClient, runtime: E2BSandboxRuntime, config: PTCSandboxConfig | None = None) -> None:
+    def __init__(self, mcp_client: BaseMCPClient | None, runtime: E2BSandboxRuntime, config: PTCSandboxConfig | None = None) -> None:
         """Initialize PTCSandboxExecutor.
 
         Args:
-            mcp_client: The MCP client with configured servers.
+            mcp_client: The MCP client with configured servers. Can be None for custom-only configs.
             runtime: The E2B sandbox runtime instance.
             config: Optional sandbox executor configuration.
 
         Raises:
             ImportError: If sandbox dependencies are not available.
         """
-    async def execute_code(self, code: str) -> SandboxExecutionResult:
+    async def execute_code(self, code: str, tool_configs: dict[str, dict] | None = None) -> SandboxExecutionResult:
         """Execute code inside the sandbox with MCP access.
 
         This method:
-        1. Builds the sandbox payload (MCP config + generated tool modules)
+        1. Builds the sandbox payload (MCP config + generated tool modules + custom tools)
         2. Wraps the user code with necessary imports and setup
         3. Executes the code in the E2B sandbox
         4. Returns the execution result (stdout/stderr/exit_code)
 
         Args:
             code: Python code to execute in the sandbox.
+            tool_configs: Optional per-tool config values for custom LangChain tools.
+                These are merged with agent defaults and written to custom_defaults.json.
 
         Returns:
             SandboxExecutionResult with stdout, stderr, and exit_code.
@@ -71,3 +91,9 @@ class PTCSandboxExecutor:
         Raises:
             PTCToolError: If sandbox execution fails.
         """
+    def reset_static_bundle(self) -> None:
+        """Reset the static bundle upload state.
+
+        Call this method when the sandbox is destroyed/recreated to force
+        a full re-upload of the static bundle on the next execution.
+        """

aip_agents/ptc/naming.py

@@ -12,7 +12,7 @@ import re
 from typing import Any
 
 # Reserved module names that cannot be used for server packages.
-RESERVED_MODULE_NAMES = frozenset({"ptc_helper", "mcp_client"})
+RESERVED_MODULE_NAMES = frozenset({"ptc_helper", "mcp_client", "custom"})
 
 # Default placeholder for example values.
 DEFAULT_EXAMPLE_PLACEHOLDER = '"example"'
@@ -87,6 +87,18 @@ def sanitize_param_name(name: str) -> str:
     return sanitize_module_name(name)
 
 
+def is_valid_identifier(name: str) -> bool:
+    """Check if a string is a valid Python identifier and not a keyword.
+
+    Args:
+        name: Candidate identifier.
+
+    Returns:
+        True if name is a valid identifier and not a keyword.
+    """
+    return name.isidentifier() and not keyword.iskeyword(name)
+
+
 def json_type_to_python(json_type: str | list) -> str:
     """Convert JSON schema type to Python type hint.
 

aip_agents/ptc/naming.pyi

@@ -43,6 +43,15 @@ def sanitize_param_name(name: str) -> str:
     Returns:
         Valid Python parameter name (lowercase, underscored).
     """
+def is_valid_identifier(name: str) -> bool:
+    """Check if a string is a valid Python identifier and not a keyword.
+
+    Args:
+        name: Candidate identifier.
+
+    Returns:
+        True if name is a valid identifier and not a keyword.
+    """
 def json_type_to_python(json_type: str | list) -> str:
     """Convert JSON schema type to Python type hint.