Add E2BSandboxToolset integration for Haystack agents

haystack_experimental/tools/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0

haystack_experimental/tools/e2b/__init__.py

@@ -0,0 +1,28 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# 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)

haystack_experimental/tools/e2b/bash_tool.py

@@ -0,0 +1,82 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# 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)

haystack_experimental/tools/e2b/e2b_agent_example.py

@@ -0,0 +1,101 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# 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)

haystack_experimental/tools/e2b/e2b_pipeline_example.py

@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# 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.")