diff --git a/haystack_experimental/tools/__init__.py b/haystack_experimental/tools/__init__.py new file mode 100644 index 00000000..c1764a6e --- /dev/null +++ b/haystack_experimental/tools/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/haystack_experimental/tools/e2b/__init__.py b/haystack_experimental/tools/e2b/__init__.py new file mode 100644 index 00000000..fa0a970b --- /dev/null +++ b/haystack_experimental/tools/e2b/__init__.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +import sys +from typing import TYPE_CHECKING + +from lazy_imports import LazyImporter + +_import_structure = { + "e2b_sandbox": ["E2BSandbox"], + "bash_tool": ["RunBashCommandTool"], + "read_file_tool": ["ReadFileTool"], + "write_file_tool": ["WriteFileTool"], + "list_directory_tool": ["ListDirectoryTool"], + "sandbox_toolset": ["E2BToolset"], +} + +if TYPE_CHECKING: + from .bash_tool import RunBashCommandTool as RunBashCommandTool + from .e2b_sandbox import E2BSandbox as E2BSandbox + from .list_directory_tool import ListDirectoryTool as ListDirectoryTool + from .read_file_tool import ReadFileTool as ReadFileTool + from .sandbox_toolset import E2BToolset as E2BToolset + from .write_file_tool import WriteFileTool as WriteFileTool + +else: + sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure) diff --git a/haystack_experimental/tools/e2b/bash_tool.py b/haystack_experimental/tools/e2b/bash_tool.py new file mode 100644 index 00000000..67f85c38 --- /dev/null +++ b/haystack_experimental/tools/e2b/bash_tool.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack.core.serialization import generate_qualified_class_name +from haystack.tools import Tool + +from haystack_experimental.tools.e2b.e2b_sandbox import E2BSandbox + + +class RunBashCommandTool(Tool): + """ + A :class:`~haystack.tools.Tool` that executes bash commands inside an E2B sandbox. + + Pass the same :class:`E2BSandbox` instance to multiple tool classes so they + all operate in the same live sandbox environment. + + ### Usage example + + ```python + from haystack_experimental.tools.e2b import E2BSandbox, RunBashCommandTool, ReadFileTool + + sandbox = E2BSandbox() + agent = Agent( + chat_generator=..., + tools=[ + RunBashCommandTool(sandbox=sandbox), + ReadFileTool(sandbox=sandbox), + ], + ) + ``` + """ + + def __init__(self, sandbox: E2BSandbox) -> None: + """ + :param sandbox: The :class:`E2BSandbox` instance that will execute commands. + """ + + def run_bash_command(command: str, timeout: int = 60) -> str: + sb = sandbox._require_sandbox() + try: + result = sb.commands.run(command, timeout=timeout) + return f"exit_code: {result.exit_code}\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}" + except Exception as e: + raise RuntimeError(f"Failed to run bash command: {e}") from e + + super().__init__( + name="run_bash_command", + description=( + "Execute a bash command inside the E2B sandbox and return the combined stdout, " + "stderr, and exit code. Use this to run shell scripts, install packages, compile " + "code, or perform any system-level operation." + ), + parameters={ + "type": "object", + "properties": { + "command": {"type": "string", "description": "The bash command to execute."}, + "timeout": { + "type": "integer", + "description": ( + "Maximum number of seconds to wait for the command to finish. Defaults to 60 seconds." + ), + }, + }, + "required": ["command"], + }, + function=run_bash_command, + ) + self._e2b_sandbox = sandbox + + def to_dict(self) -> dict[str, Any]: + return { + "type": generate_qualified_class_name(type(self)), + "data": {"sandbox": self._e2b_sandbox.to_dict()}, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "RunBashCommandTool": + sandbox = E2BSandbox.from_dict(data["data"]["sandbox"]) + return cls(sandbox=sandbox) diff --git a/haystack_experimental/tools/e2b/e2b_agent_example.py b/haystack_experimental/tools/e2b/e2b_agent_example.py new file mode 100644 index 00000000..85436860 --- /dev/null +++ b/haystack_experimental/tools/e2b/e2b_agent_example.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Example: Haystack Agent with E2B sandbox tools. + +Demonstrates that all four tools (run_bash_command, read_file, write_file, +list_directory) share the same sandbox instance, so the agent can write a +file in one step and read it back / execute it in the next. + +Requirements: + pip install haystack-ai e2b openai + +Environment variables: + E2B_API_KEY – your E2B API key + OPENAI_API_KEY – your OpenAI API key (or swap the generator below) +""" + +import sys + +from haystack.components.generators.chat import OpenAIChatGenerator +from haystack.dataclasses import ChatMessage +from haystack.utils import Secret +from haystack_experimental.components.agents import Agent +from haystack_experimental.tools.e2b import ( + E2BSandbox, + ListDirectoryTool, + ReadFileTool, + RunBashCommandTool, + WriteFileTool, +) + +# --------------------------------------------------------------------------- +# Example queries that exercise cross-tool state sharing: +# 1. The agent writes a Python script to the sandbox filesystem. +# 2. It executes the script via bash and captures stdout. +# 3. It reads the output file back (or lists a directory) to verify results. +# --------------------------------------------------------------------------- +EXAMPLE_QUERIES = [ + # Simple: purely bash-based data wrangling + ( + "Generate the first 10 Fibonacci numbers using a bash one-liner " + "and show me the results." + ), + # Cross-tool: write → execute → read + ( + "Write a Python script to /tmp/primes.py that prints all prime numbers " + "up to 50, run it, and then read the file back so I can see both the " + "script and its output." + ), + # Multi-step: write → list → bash + ( + "Create a directory /tmp/workspace, write three small text files into it " + "with different content, list the directory to confirm they exist, and " + "then use bash to count the total number of words across all three files." + ), +] + + +def run(query: str, model: str = "gpt-4o-mini") -> None: + print("\n" + "=" * 70) + print(f"Query: {query}") + print("=" * 70) + + # One sandbox passed to each tool class – they all share the same live sandbox process. + sandbox = E2BSandbox() + sandbox.warm_up() + tools = [ + RunBashCommandTool(sandbox=sandbox), + ReadFileTool(sandbox=sandbox), + WriteFileTool(sandbox=sandbox), + ListDirectoryTool(sandbox=sandbox), + ] + + agent = Agent( + chat_generator=OpenAIChatGenerator(model=model), + tools=tools, + system_prompt=( + "You are a helpful coding assistant with access to a live Linux sandbox. " + "Use the available tools freely to explore, write files, and run commands. " + "All tools operate inside the same sandbox environment, so files written " + "with write_file are immediately available to run_bash_command and read_file." + ), + max_agent_steps=15, + ) + + result = agent.run(messages=[ChatMessage.from_user(query)]) + print("\n--- Agent response ---") + print(result["last_message"].text) + + + +if __name__ == "__main__": + # Run a specific query index (0/1/2) or all of them by default. + if len(sys.argv) > 1: + idx = int(sys.argv[1]) + run(EXAMPLE_QUERIES[idx]) + else: + for query in EXAMPLE_QUERIES: + run(query) diff --git a/haystack_experimental/tools/e2b/e2b_pipeline_example.py b/haystack_experimental/tools/e2b/e2b_pipeline_example.py new file mode 100644 index 00000000..27c9b999 --- /dev/null +++ b/haystack_experimental/tools/e2b/e2b_pipeline_example.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Example: Haystack Pipeline with an Agent and E2B sandbox tools. + +Demonstrates that a Pipeline containing an Agent with E2BToolset can be: + 1. Serialised to YAML + 2. Written to disk + 3. Loaded back from YAML with full sandbox config intact + +All four tools (run_bash_command, read_file, write_file, list_directory) +share a single E2BSandbox after the round-trip, so the agent operates +in one live sandbox environment. + +Requirements: + pip install haystack-ai e2b openai + +Environment variables: + E2B_API_KEY – your E2B API key + OPENAI_API_KEY – your OpenAI API key +""" + +import tempfile +from pathlib import Path + +from haystack.components.generators.chat import OpenAIChatGenerator +from haystack.core.pipeline import Pipeline +from haystack.dataclasses import ChatMessage + +from haystack_experimental.components.agents import Agent +from haystack_experimental.tools.e2b import E2BToolset + + +def build_pipeline() -> Pipeline: + agent = Agent( + chat_generator=OpenAIChatGenerator(model="gpt-4o-mini"), + tools=E2BToolset(sandbox_template="base", timeout=120), + system_prompt=( + "You are a helpful coding assistant with access to a live Linux sandbox. " + "Use the available tools freely to explore, write files, and run commands." + ), + max_agent_steps=10, + ) + pipeline = Pipeline() + pipeline.add_component("agent", agent) + return pipeline + + +def roundtrip_yaml(pipeline: Pipeline) -> Pipeline: + """Serialise to YAML, save to a temp file, load it back.""" + yaml_str = pipeline.dumps() + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write(yaml_str) + yaml_path = Path(f.name) + + print(f"Pipeline YAML written to {yaml_path}\n") + print(yaml_str) + print("---\n") + + restored = Pipeline.loads(yaml_path.read_text()) + return restored + + +def verify_roundtrip(original: Pipeline, restored: Pipeline) -> None: + """Check that the restored pipeline has the same structure.""" + orig_agent: Agent = original.get_component("agent") + rest_agent: Agent = restored.get_component("agent") + + orig_ts: E2BToolset = orig_agent.tools # type: ignore[assignment] + rest_ts: E2BToolset = rest_agent.tools # type: ignore[assignment] + + assert type(rest_ts).__name__ == "E2BToolset", "Toolset type mismatch" + assert [t.name for t in rest_ts] == [t.name for t in orig_ts], "Tool names mismatch" + assert rest_ts.sandbox.sandbox_template == orig_ts.sandbox.sandbox_template + assert rest_ts.sandbox.timeout == orig_ts.sandbox.timeout + + sandbox_ids = {id(t._e2b_sandbox) for t in rest_ts} + assert len(sandbox_ids) == 1, "Tools should share a single sandbox after round-trip" + + print("All assertions passed: YAML round-trip preserves pipeline structure.\n") + + +def run_agent(pipeline: Pipeline, query: str) -> None: + """Run the agent with a query (requires live API keys).""" + print(f"Query: {query}\n") + result = pipeline.run(data={"agent": {"messages": [ChatMessage.from_user(query)]}}) + print("--- Agent response ---") + print(result["agent"]["last_message"].text) + + +if __name__ == "__main__": + pipeline = build_pipeline() + restored = roundtrip_yaml(pipeline) + verify_roundtrip(pipeline, restored) + + run_agent(restored, "Write a Python one-liner to /tmp/hello.py that prints 'Hello from E2B!', run it, then show me the output.") diff --git a/haystack_experimental/tools/e2b/e2b_sandbox.py b/haystack_experimental/tools/e2b/e2b_sandbox.py new file mode 100644 index 00000000..1693a81c --- /dev/null +++ b/haystack_experimental/tools/e2b/e2b_sandbox.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack import logging +from haystack.lazy_imports import LazyImport +from haystack.utils import Secret, deserialize_secrets_inplace + +with LazyImport(message="Run 'pip install e2b'") as e2b_import: + from e2b import Sandbox + +logger = logging.getLogger(__name__) + + +class E2BSandbox: + """ + Manages the lifecycle of an E2B cloud sandbox. + + Instantiate this class and pass it to one or more E2B tool classes + (``RunBashCommandTool``, ``ReadFileTool``, ``WriteFileTool``, + ``ListDirectoryTool``) to share a single sandbox environment across all + tools. All tools that receive the same ``E2BSandbox`` instance operate + inside the same live sandbox process. + + ### Usage example + + ```python + from haystack.components.generators.chat import OpenAIChatGenerator + + from haystack_experimental.components.agents import Agent + from haystack_experimental.tools.e2b import ( + E2BSandbox, + RunBashCommandTool, + ReadFileTool, + WriteFileTool, + ListDirectoryTool, + ) + + sandbox = E2BSandbox() + agent = Agent( + chat_generator=OpenAIChatGenerator(model="gpt-4o"), + tools=[ + RunBashCommandTool(sandbox=sandbox), + ReadFileTool(sandbox=sandbox), + WriteFileTool(sandbox=sandbox), + ListDirectoryTool(sandbox=sandbox), + ], + ) + ``` + + Lifecycle is handled automatically by the Agent's pipeline. If you use the + tools standalone, call :meth:`warm_up` before the first tool invocation: + + ```python + sandbox.warm_up() + # … use tools … + sandbox.close() + ``` + """ + + def __init__( + self, + api_key: Secret | None = None, + sandbox_template: str = "base", + timeout: int = 120, + environment_vars: dict[str, str] | None = None, + ): + self.api_key = api_key or Secret.from_env_var("E2B_API_KEY") + self.sandbox_template = sandbox_template + self.timeout = timeout + self.environment_vars = environment_vars or {} + self._sandbox: Any = None + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def warm_up(self) -> None: + """ + Establish the connection to the E2B sandbox. + + Idempotent – calling it multiple times has no effect if the sandbox is + already running. + + :raises RuntimeError: If the E2B sandbox cannot be created. + """ + if self._sandbox is not None: + return + + e2b_import.check() + resolved_key = self.api_key.resolve_value() + try: + logger.info( + "Starting E2B sandbox (template={template}, timeout={timeout}s)", + template=self.sandbox_template, + timeout=self.timeout, + ) + self._sandbox = Sandbox.create( + api_key=resolved_key, + template=self.sandbox_template, + timeout=self.timeout, + envs=self.environment_vars if self.environment_vars else None, + ) + logger.info("E2B sandbox started (id={sandbox_id})", sandbox_id=self._sandbox.sandbox_id) + except Exception as e: + raise RuntimeError(f"Failed to start E2B sandbox: {e}") from e + + def close(self) -> None: + """ + Shut down the E2B sandbox and release all associated resources. + + Call this when you are done to avoid leaving idle sandboxes running. + """ + if self._sandbox is None: + return + try: + self._sandbox.kill() + logger.info("E2B sandbox closed") + except Exception as e: + logger.warning("Failed to close E2B sandbox: {error}", error=e) + finally: + self._sandbox = None + + # ------------------------------------------------------------------ + # Serialisation + # ------------------------------------------------------------------ + + def to_dict(self) -> dict[str, Any]: + """ + Serialize the sandbox configuration to a dictionary. + + :returns: Dictionary containing the serialised configuration. + """ + from haystack.core.serialization import generate_qualified_class_name + + return { + "type": generate_qualified_class_name(type(self)), + "data": { + "api_key": self.api_key.to_dict(), + "sandbox_template": self.sandbox_template, + "timeout": self.timeout, + "environment_vars": self.environment_vars, + }, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "E2BSandbox": + """ + Deserialize an :class:`E2BSandbox` from a dictionary. + + :param data: Dictionary created by :meth:`to_dict`. + :returns: A new :class:`E2BSandbox` instance ready to be warmed up. + """ + inner = data["data"] + deserialize_secrets_inplace(inner, keys=["api_key"]) + return cls( + api_key=inner["api_key"], + sandbox_template=inner.get("sandbox_template", "base"), + timeout=inner.get("timeout", 120), + environment_vars=inner.get("environment_vars", {}), + ) + + # ------------------------------------------------------------------ + # Internal helpers (used by the tool classes) + # ------------------------------------------------------------------ + + def _require_sandbox(self) -> "Sandbox": + """Return the active sandbox or raise a helpful error.""" + if self._sandbox is None: + raise RuntimeError( + "E2B sandbox is not running. Call warm_up() before using the tools, " + "or add the sandbox to a Haystack pipeline/agent which calls warm_up() automatically." + ) + return self._sandbox diff --git a/haystack_experimental/tools/e2b/list_directory_tool.py b/haystack_experimental/tools/e2b/list_directory_tool.py new file mode 100644 index 00000000..38d5d114 --- /dev/null +++ b/haystack_experimental/tools/e2b/list_directory_tool.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack.core.serialization import generate_qualified_class_name +from haystack.tools import Tool + +from haystack_experimental.tools.e2b.e2b_sandbox import E2BSandbox + + +class ListDirectoryTool(Tool): + """ + A :class:`~haystack.tools.Tool` that lists directory contents in an E2B sandbox. + + Pass the same :class:`E2BSandbox` instance to multiple tool classes so they + all operate in the same live sandbox environment. + + ### Usage example + + ```python + from haystack_experimental.tools.e2b import E2BSandbox, ListDirectoryTool + + sandbox = E2BSandbox() + agent = Agent(chat_generator=..., tools=[ListDirectoryTool(sandbox=sandbox)]) + ``` + """ + + def __init__(self, sandbox: E2BSandbox) -> None: + """ + :param sandbox: The :class:`E2BSandbox` instance to list directories from. + """ + + def list_directory(path: str) -> str: + sb = sandbox._require_sandbox() + try: + entries = sb.files.list(path) + lines = [] + for entry in entries: + name = entry.name + if getattr(entry, "is_dir", False) or getattr(entry, "type", "") == "dir": + name = name + "/" + lines.append(name) + return "\n".join(lines) if lines else "(empty directory)" + except Exception as e: + raise RuntimeError(f"Failed to list directory '{path}': {e}") from e + + super().__init__( + name="list_directory", + description=( + "List the files and subdirectories inside a directory in the E2B sandbox " + "filesystem. Returns a newline-separated list of names with a trailing '/' " + "appended to subdirectory names." + ), + parameters={ + "type": "object", + "properties": { + "path": {"type": "string", "description": "Absolute or relative path of the directory to list."} + }, + "required": ["path"], + }, + function=list_directory, + ) + self._e2b_sandbox = sandbox + + def to_dict(self) -> dict[str, Any]: + return { + "type": generate_qualified_class_name(type(self)), + "data": {"sandbox": self._e2b_sandbox.to_dict()}, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "ListDirectoryTool": + sandbox = E2BSandbox.from_dict(data["data"]["sandbox"]) + return cls(sandbox=sandbox) diff --git a/haystack_experimental/tools/e2b/read_file_tool.py b/haystack_experimental/tools/e2b/read_file_tool.py new file mode 100644 index 00000000..f2e676d4 --- /dev/null +++ b/haystack_experimental/tools/e2b/read_file_tool.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack.core.serialization import generate_qualified_class_name +from haystack.tools import Tool + +from haystack_experimental.tools.e2b.e2b_sandbox import E2BSandbox + + +class ReadFileTool(Tool): + """ + A :class:`~haystack.tools.Tool` that reads files from an E2B sandbox filesystem. + + Pass the same :class:`E2BSandbox` instance to multiple tool classes so they + all operate in the same live sandbox environment. + + ### Usage example + + ```python + from haystack_experimental.tools.e2b import E2BSandbox, ReadFileTool + + sandbox = E2BSandbox() + agent = Agent(chat_generator=..., tools=[ReadFileTool(sandbox=sandbox)]) + ``` + """ + + def __init__(self, sandbox: E2BSandbox) -> None: + """ + :param sandbox: The :class:`E2BSandbox` instance to read files from. + """ + + def read_file(path: str) -> str: + sb = sandbox._require_sandbox() + try: + content = sb.files.read(path) + if isinstance(content, bytes): + return content.decode("utf-8", errors="replace") + return str(content) + except Exception as e: + raise RuntimeError(f"Failed to read file '{path}': {e}") from e + + super().__init__( + name="read_file", + description=( + "Read the text content of a file from the E2B sandbox filesystem and return it " + "as a string. The file must exist; use list_directory to verify paths first." + ), + parameters={ + "type": "object", + "properties": { + "path": {"type": "string", "description": "Absolute or relative path of the file to read."} + }, + "required": ["path"], + }, + function=read_file, + ) + self._e2b_sandbox = sandbox + + def to_dict(self) -> dict[str, Any]: + return { + "type": generate_qualified_class_name(type(self)), + "data": {"sandbox": self._e2b_sandbox.to_dict()}, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "ReadFileTool": + sandbox = E2BSandbox.from_dict(data["data"]["sandbox"]) + return cls(sandbox=sandbox) diff --git a/haystack_experimental/tools/e2b/sandbox_toolset.py b/haystack_experimental/tools/e2b/sandbox_toolset.py new file mode 100644 index 00000000..04344efc --- /dev/null +++ b/haystack_experimental/tools/e2b/sandbox_toolset.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack.core.serialization import generate_qualified_class_name +from haystack.tools import Toolset +from haystack.utils import Secret + +from haystack_experimental.tools.e2b.bash_tool import RunBashCommandTool +from haystack_experimental.tools.e2b.e2b_sandbox import E2BSandbox +from haystack_experimental.tools.e2b.list_directory_tool import ListDirectoryTool +from haystack_experimental.tools.e2b.read_file_tool import ReadFileTool +from haystack_experimental.tools.e2b.write_file_tool import WriteFileTool + + +class E2BToolset(Toolset): + """ + A :class:`~haystack.tools.Toolset` that bundles all E2B sandbox tools. + + All tools in the set share a single :class:`E2BSandbox` instance so they + operate inside the same live sandbox process. The toolset owns the sandbox + lifecycle: calling :meth:`warm_up` starts the sandbox, and serialisation + round-trips preserve the shared-sandbox relationship. + + ### Usage example + + ```python + from haystack.components.generators.chat import OpenAIChatGenerator + + from haystack_experimental.components.agents import Agent + from haystack_experimental.tools.e2b import E2BToolset + + agent = Agent( + chat_generator=OpenAIChatGenerator(model="gpt-4o"), + tools=E2BToolset(), + ) + ``` + """ + + def __init__( + self, + api_key: Secret | None = None, + sandbox_template: str = "base", + timeout: int = 120, + environment_vars: dict[str, str] | None = None, + ): + """ + :param api_key: E2B API key. Defaults to ``Secret.from_env_var("E2B_API_KEY")``. + :param sandbox_template: E2B sandbox template name. Defaults to ``"base"``. + :param timeout: Sandbox inactivity timeout in seconds. Defaults to ``300``. + :param environment_vars: Optional environment variables to inject into the sandbox. + """ + self.sandbox = E2BSandbox( + api_key=api_key, + sandbox_template=sandbox_template, + timeout=timeout, + environment_vars=environment_vars, + ) + super().__init__( + tools=[ + RunBashCommandTool(sandbox=self.sandbox), + ReadFileTool(sandbox=self.sandbox), + WriteFileTool(sandbox=self.sandbox), + ListDirectoryTool(sandbox=self.sandbox), + ] + ) + + def warm_up(self) -> None: + """Start the shared E2B sandbox (idempotent).""" + self.sandbox.warm_up() + + def close(self) -> None: + """Shut down the shared E2B sandbox and release cloud resources.""" + self.sandbox.close() + + def to_dict(self) -> dict[str, Any]: + return { + "type": generate_qualified_class_name(type(self)), + "data": self.sandbox.to_dict()["data"], + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "E2BToolset": + from haystack.utils import deserialize_secrets_inplace + + inner = data["data"] + deserialize_secrets_inplace(inner, keys=["api_key"]) + return cls( + api_key=inner["api_key"], + sandbox_template=inner.get("sandbox_template", "base"), + timeout=inner.get("timeout", 120), + environment_vars=inner.get("environment_vars", {}), + ) diff --git a/haystack_experimental/tools/e2b/write_file_tool.py b/haystack_experimental/tools/e2b/write_file_tool.py new file mode 100644 index 00000000..304d2b95 --- /dev/null +++ b/haystack_experimental/tools/e2b/write_file_tool.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from haystack.core.serialization import generate_qualified_class_name +from haystack.tools import Tool + +from haystack_experimental.tools.e2b.e2b_sandbox import E2BSandbox + + +class WriteFileTool(Tool): + """ + A :class:`~haystack.tools.Tool` that writes files to an E2B sandbox filesystem. + + Pass the same :class:`E2BSandbox` instance to multiple tool classes so they + all operate in the same live sandbox environment. + + ### Usage example + + ```python + from haystack_experimental.tools.e2b import E2BSandbox, WriteFileTool + + sandbox = E2BSandbox() + agent = Agent(chat_generator=..., tools=[WriteFileTool(sandbox=sandbox)]) + ``` + """ + + def __init__(self, sandbox: E2BSandbox) -> None: + """ + :param sandbox: The :class:`E2BSandbox` instance to write files to. + """ + + def write_file(path: str, content: str) -> str: + sb = sandbox._require_sandbox() + try: + sb.files.write(path, content) + return f"File written successfully: {path}" + except Exception as e: + raise RuntimeError(f"Failed to write file '{path}': {e}") from e + + super().__init__( + name="write_file", + description=( + "Write text content to a file in the E2B sandbox filesystem. " + "Parent directories are created automatically if they do not exist. " + "Existing files are overwritten." + ), + parameters={ + "type": "object", + "properties": { + "path": {"type": "string", "description": "Absolute or relative path of the file to write."}, + "content": {"type": "string", "description": "Text content to write into the file."}, + }, + "required": ["path", "content"], + }, + function=write_file, + ) + self._e2b_sandbox = sandbox + + def to_dict(self) -> dict[str, Any]: + return { + "type": generate_qualified_class_name(type(self)), + "data": {"sandbox": self._e2b_sandbox.to_dict()}, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "WriteFileTool": + sandbox = E2BSandbox.from_dict(data["data"]["sandbox"]) + return cls(sandbox=sandbox) diff --git a/pyproject.toml b/pyproject.toml index d91ccfdd..624d56db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ extra-dependencies = [ "tiktoken", # LLM-based Summarizer "nltk>=3.9.1", # LLM-based Summarizer "mem0ai", # Mem0MemoryStore + "e2b", # E2BSandboxToolset "amazon-bedrock-haystack", "google-genai-haystack", "cohere-haystack", diff --git a/test/tools/__init__.py b/test/tools/__init__.py new file mode 100644 index 00000000..c1764a6e --- /dev/null +++ b/test/tools/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/test/tools/e2b/__init__.py b/test/tools/e2b/__init__.py new file mode 100644 index 00000000..c1764a6e --- /dev/null +++ b/test/tools/e2b/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 diff --git a/test/tools/e2b/test_sandbox_toolset.py b/test/tools/e2b/test_sandbox_toolset.py new file mode 100644 index 00000000..84f4f728 --- /dev/null +++ b/test/tools/e2b/test_sandbox_toolset.py @@ -0,0 +1,434 @@ +# SPDX-FileCopyrightText: 2022-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import MagicMock, patch + +import pytest +from haystack.tools.errors import ToolInvocationError +from haystack.utils import Secret + +from haystack_experimental.tools.e2b.bash_tool import RunBashCommandTool +from haystack_experimental.tools.e2b.e2b_sandbox import E2BSandbox +from haystack_experimental.tools.e2b.list_directory_tool import ListDirectoryTool +from haystack_experimental.tools.e2b.read_file_tool import ReadFileTool +from haystack_experimental.tools.e2b.sandbox_toolset import E2BToolset +from haystack_experimental.tools.e2b.write_file_tool import WriteFileTool + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_sandbox(**kwargs) -> E2BSandbox: + """Create an E2BSandbox with a dummy API key for testing.""" + defaults = {"api_key": Secret.from_token("test-api-key")} + defaults.update(kwargs) + return E2BSandbox(**defaults) + + +def _make_sandbox_mock() -> MagicMock: + """Return a MagicMock that mimics the e2b Sandbox object.""" + sandbox = MagicMock() + sandbox.sandbox_id = "sandbox-test-123" + return sandbox + + +def _sandbox_with_mock() -> tuple[E2BSandbox, MagicMock]: + """Return an E2BSandbox that already has a mocked underlying sandbox.""" + sb = _make_sandbox() + mock = _make_sandbox_mock() + sb._sandbox = mock + return sb, mock + + +# --------------------------------------------------------------------------- +# E2BSandbox – initialisation +# --------------------------------------------------------------------------- + + +class TestE2BSandboxInit: + def test_class_defaults(self): + """Verify the real class defaults, not values set by a helper.""" + sandbox = E2BSandbox(api_key=Secret.from_token("test-api-key")) + assert sandbox.sandbox_template == "base" + assert sandbox.timeout == 120 + assert sandbox.environment_vars == {} + assert sandbox._sandbox is None + + def test_custom_parameters(self): + sandbox = _make_sandbox( + sandbox_template="my-template", + timeout=600, + environment_vars={"FOO": "bar"}, + ) + assert sandbox.sandbox_template == "my-template" + assert sandbox.timeout == 600 + assert sandbox.environment_vars == {"FOO": "bar"} + + +# --------------------------------------------------------------------------- +# E2BSandbox – warm_up +# --------------------------------------------------------------------------- + + +class TestE2BSandboxWarmUp: + @patch("haystack_experimental.tools.e2b.e2b_sandbox.e2b_import") + @patch("haystack_experimental.tools.e2b.e2b_sandbox.Sandbox.create") + def test_warm_up_creates_sandbox(self, mock_sandbox_create, mock_e2b_import): + mock_e2b_import.check.return_value = None + mock_instance = _make_sandbox_mock() + mock_sandbox_create.return_value = mock_instance + + sb = _make_sandbox(sandbox_template="base", timeout=120) + sb.warm_up() + + mock_sandbox_create.assert_called_once_with( + api_key="test-api-key", + template="base", + timeout=120, + envs=None, + ) + assert sb._sandbox is mock_instance + + @patch("haystack_experimental.tools.e2b.e2b_sandbox.e2b_import") + @patch("haystack_experimental.tools.e2b.e2b_sandbox.Sandbox.create") + def test_warm_up_passes_environment_vars(self, mock_sandbox_create, mock_e2b_import): + mock_e2b_import.check.return_value = None + mock_sandbox_create.return_value = _make_sandbox_mock() + + sb = _make_sandbox(environment_vars={"MY_VAR": "value"}) + sb.warm_up() + + _, kwargs = mock_sandbox_create.call_args + assert kwargs["envs"] == {"MY_VAR": "value"} + + @patch("haystack_experimental.tools.e2b.e2b_sandbox.e2b_import") + @patch("haystack_experimental.tools.e2b.e2b_sandbox.Sandbox.create") + def test_warm_up_is_idempotent(self, mock_sandbox_create, mock_e2b_import): + mock_e2b_import.check.return_value = None + mock_sandbox_create.return_value = _make_sandbox_mock() + + sb = _make_sandbox() + sb.warm_up() + sb.warm_up() + + mock_sandbox_create.assert_called_once() + + @patch("haystack_experimental.tools.e2b.e2b_sandbox.e2b_import") + @patch("haystack_experimental.tools.e2b.e2b_sandbox.Sandbox.create") + def test_warm_up_raises_on_sandbox_error(self, mock_sandbox_create, mock_e2b_import): + mock_e2b_import.check.return_value = None + mock_sandbox_create.side_effect = Exception("connection refused") + + sb = _make_sandbox() + with pytest.raises(RuntimeError, match="Failed to start E2B sandbox"): + sb.warm_up() + + +# --------------------------------------------------------------------------- +# E2BSandbox – close +# --------------------------------------------------------------------------- + + +class TestE2BSandboxClose: + def test_close_without_warm_up_is_noop(self): + sb = _make_sandbox() + sb.close() + assert sb._sandbox is None + + def test_close_kills_sandbox(self): + sb, mock = _sandbox_with_mock() + sb.close() + mock.kill.assert_called_once() + assert sb._sandbox is None + + def test_close_clears_sandbox_on_kill_error(self): + sb, mock = _sandbox_with_mock() + mock.kill.side_effect = Exception("kill failed") + sb.close() # must not raise + assert sb._sandbox is None + + +# --------------------------------------------------------------------------- +# E2BSandbox – serialisation +# --------------------------------------------------------------------------- + + +class TestE2BSandboxSerialisation: + def _make_env_sandbox(self, **kwargs) -> E2BSandbox: + defaults = {"api_key": Secret.from_env_var("E2B_API_KEY")} + defaults.update(kwargs) + return E2BSandbox(**defaults) + + def test_to_dict_contains_expected_keys(self): + sb = self._make_env_sandbox(sandbox_template="my-template", timeout=600) + data = sb.to_dict() + + assert "type" in data + assert "data" in data + assert data["data"]["sandbox_template"] == "my-template" + assert data["data"]["timeout"] == 600 + + def test_to_dict_does_not_include_sandbox_instance(self): + sb = self._make_env_sandbox() + sb._sandbox = _make_sandbox_mock() + data = sb.to_dict() + + assert "_sandbox" not in data["data"] + assert "sandbox" not in data["data"] + + def test_from_dict_round_trip(self): + original = self._make_env_sandbox( + sandbox_template="custom", + timeout=900, + environment_vars={"KEY": "value"}, + ) + data = original.to_dict() + restored = E2BSandbox.from_dict(data) + + assert restored.sandbox_template == "custom" + assert restored.timeout == 900 + assert restored.environment_vars == {"KEY": "value"} + assert restored._sandbox is None + + def test_to_dict_type_is_qualified_class_name(self): + sb = self._make_env_sandbox() + data = sb.to_dict() + assert "E2BSandbox" in data["type"] + + +# --------------------------------------------------------------------------- +# Tool classes – structure +# --------------------------------------------------------------------------- + + +class TestToolClasses: + def test_run_bash_command_tool_name_and_schema(self): + sb = _make_sandbox() + tool = RunBashCommandTool(sandbox=sb) + assert tool.name == "run_bash_command" + assert tool.description + assert "command" in tool.parameters["required"] + + def test_read_file_tool_name_and_schema(self): + sb = _make_sandbox() + tool = ReadFileTool(sandbox=sb) + assert tool.name == "read_file" + assert tool.description + assert "path" in tool.parameters["required"] + + def test_write_file_tool_name_and_schema(self): + sb = _make_sandbox() + tool = WriteFileTool(sandbox=sb) + assert tool.name == "write_file" + assert tool.description + assert "path" in tool.parameters["required"] + assert "content" in tool.parameters["required"] + + def test_list_directory_tool_name_and_schema(self): + sb = _make_sandbox() + tool = ListDirectoryTool(sandbox=sb) + assert tool.name == "list_directory" + assert tool.description + assert "path" in tool.parameters["required"] + + def test_tool_stores_sandbox_reference(self): + sb = _make_sandbox() + tool = RunBashCommandTool(sandbox=sb) + assert tool._e2b_sandbox is sb + + def test_e2b_toolset_contains_four_tools(self): + ts = E2BToolset(api_key=Secret.from_token("test-api-key")) + assert len(ts) == 4 + names = {t.name for t in ts} + assert names == {"run_bash_command", "read_file", "write_file", "list_directory"} + + def test_e2b_toolset_has_correct_tool_types(self): + ts = E2BToolset(api_key=Secret.from_token("test-api-key")) + tool_types = {type(t) for t in ts} + assert tool_types == {RunBashCommandTool, ReadFileTool, WriteFileTool, ListDirectoryTool} + + def test_e2b_toolset_shares_same_sandbox(self): + ts = E2BToolset(api_key=Secret.from_token("test-api-key")) + assert all(t._e2b_sandbox is ts.sandbox for t in ts) + + mock = _make_sandbox_mock() + mock.commands.run.return_value = MagicMock(exit_code=0, stdout="ok", stderr="") + ts.sandbox._sandbox = mock + + bash_tool = next(t for t in ts if t.name == "run_bash_command") + bash_tool.invoke(command="echo ok") + + mock.commands.run.assert_called_once() + + def test_e2b_toolset_default_api_key(self): + """E2BToolset uses E2B_API_KEY env var when api_key is omitted.""" + ts = E2BToolset() + assert ts.sandbox.api_key is not None + + def test_tools_from_same_sandbox_share_state(self): + """Tools instantiated with the same sandbox share state.""" + sb = _make_sandbox() + bash_tool = RunBashCommandTool(sandbox=sb) + read_tool = ReadFileTool(sandbox=sb) + assert bash_tool._e2b_sandbox is read_tool._e2b_sandbox + + +# --------------------------------------------------------------------------- +# RunBashCommandTool behaviour +# --------------------------------------------------------------------------- + + +class TestRunBashCommandTool: + def test_returns_formatted_output(self): + sb, mock = _sandbox_with_mock() + mock_result = MagicMock(exit_code=0, stdout="hello world\n", stderr="") + mock.commands.run.return_value = mock_result + tool = RunBashCommandTool(sandbox=sb) + + output = tool.invoke(command="echo hello world") + + assert "exit_code: 0" in output + assert "hello world" in output + mock.commands.run.assert_called_once_with("echo hello world", timeout=60) + + def test_passes_custom_timeout(self): + sb, mock = _sandbox_with_mock() + mock.commands.run.return_value = MagicMock(exit_code=0, stdout="", stderr="") + tool = RunBashCommandTool(sandbox=sb) + + tool.invoke(command="sleep 5", timeout=30) + + mock.commands.run.assert_called_once_with("sleep 5", timeout=30) + + def test_raises_when_no_sandbox(self): + sb = _make_sandbox() + tool = RunBashCommandTool(sandbox=sb) + with pytest.raises(ToolInvocationError, match="E2B sandbox is not running"): + tool.invoke(command="ls") + + def test_wraps_sandbox_exception(self): + sb, mock = _sandbox_with_mock() + mock.commands.run.side_effect = Exception("timeout") + tool = RunBashCommandTool(sandbox=sb) + with pytest.raises(ToolInvocationError, match="Failed to run bash command"): + tool.invoke(command="sleep 1000") + + +# --------------------------------------------------------------------------- +# ReadFileTool behaviour +# --------------------------------------------------------------------------- + + +class TestReadFileTool: + def test_returns_string(self): + sb, mock = _sandbox_with_mock() + mock.files.read.return_value = "file content" + tool = ReadFileTool(sandbox=sb) + + result = tool.invoke(path="/some/file.txt") + + assert result == "file content" + mock.files.read.assert_called_once_with("/some/file.txt") + + def test_decodes_bytes(self): + sb, mock = _sandbox_with_mock() + mock.files.read.return_value = b"binary content" + tool = ReadFileTool(sandbox=sb) + + result = tool.invoke(path="/binary.bin") + + assert result == "binary content" + + def test_raises_when_no_sandbox(self): + sb = _make_sandbox() + tool = ReadFileTool(sandbox=sb) + with pytest.raises(ToolInvocationError, match="E2B sandbox is not running"): + tool.invoke(path="/some/file.txt") + + def test_wraps_sandbox_exception(self): + sb, mock = _sandbox_with_mock() + mock.files.read.side_effect = Exception("file not found") + tool = ReadFileTool(sandbox=sb) + with pytest.raises(ToolInvocationError, match="Failed to read file"): + tool.invoke(path="/nonexistent.txt") + + +# --------------------------------------------------------------------------- +# WriteFileTool behaviour +# --------------------------------------------------------------------------- + + +class TestWriteFileTool: + def test_returns_confirmation(self): + sb, mock = _sandbox_with_mock() + tool = WriteFileTool(sandbox=sb) + + result = tool.invoke(path="/output/result.txt", content="hello") + + assert "/output/result.txt" in result + mock.files.write.assert_called_once_with("/output/result.txt", "hello") + + def test_raises_when_no_sandbox(self): + sb = _make_sandbox() + tool = WriteFileTool(sandbox=sb) + with pytest.raises(ToolInvocationError, match="E2B sandbox is not running"): + tool.invoke(path="/some/path.txt", content="content") + + def test_wraps_sandbox_exception(self): + sb, mock = _sandbox_with_mock() + mock.files.write.side_effect = Exception("permission denied") + tool = WriteFileTool(sandbox=sb) + with pytest.raises(ToolInvocationError, match="Failed to write file"): + tool.invoke(path="/protected/file.txt", content="data") + + +# --------------------------------------------------------------------------- +# ListDirectoryTool behaviour +# --------------------------------------------------------------------------- + + +class TestListDirectoryTool: + def _make_entry(self, name: str, is_dir: bool = False) -> MagicMock: + entry = MagicMock() + entry.name = name + entry.is_dir = is_dir + return entry + + def test_returns_names(self): + sb, mock = _sandbox_with_mock() + mock.files.list.return_value = [ + self._make_entry("file.txt"), + self._make_entry("subdir", is_dir=True), + ] + tool = ListDirectoryTool(sandbox=sb) + + result = tool.invoke(path="/home/user") + + assert "file.txt" in result + assert "subdir/" in result + mock.files.list.assert_called_once_with("/home/user") + + def test_empty_directory(self): + sb, mock = _sandbox_with_mock() + mock.files.list.return_value = [] + tool = ListDirectoryTool(sandbox=sb) + + result = tool.invoke(path="/empty") + + assert result == "(empty directory)" + + def test_raises_when_no_sandbox(self): + sb = _make_sandbox() + tool = ListDirectoryTool(sandbox=sb) + with pytest.raises(ToolInvocationError, match="E2B sandbox is not running"): + tool.invoke(path="/home") + + def test_wraps_sandbox_exception(self): + sb, mock = _sandbox_with_mock() + mock.files.list.side_effect = Exception("not a directory") + tool = ListDirectoryTool(sandbox=sb) + with pytest.raises(ToolInvocationError, match="Failed to list directory"): + tool.invoke(path="/nonexistent")