From f3470f4bde5d2b801d82d91f83dac020beab1b4b Mon Sep 17 00:00:00 2001
From: Matthew Zhou <mattzh1314@gmail.com>
Date: Wed, 4 Dec 2024 15:20:57 -0800
Subject: [PATCH] fix: Suppress warnings during local sandbox tool execution w/
 venv (#2164)

---
 letta/services/tool_execution_sandbox.py      |  2 ++
 ...integration_test_tool_execution_sandbox.py | 36 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/letta/services/tool_execution_sandbox.py b/letta/services/tool_execution_sandbox.py
index 4dedfb3b..98c8603a 100644
--- a/letta/services/tool_execution_sandbox.py
+++ b/letta/services/tool_execution_sandbox.py
@@ -159,6 +159,8 @@ class ToolExecutionSandbox:
         # Set up env for venv
         env["VIRTUAL_ENV"] = venv_path
         env["PATH"] = os.path.join(venv_path, "bin") + ":" + env["PATH"]
+        # Suppress all warnings
+        env["PYTHONWARNINGS"] = "ignore"
 
         # Execute the code in a restricted subprocess
         try:
diff --git a/tests/integration_test_tool_execution_sandbox.py b/tests/integration_test_tool_execution_sandbox.py
index bb414fae..09129e49 100644
--- a/tests/integration_test_tool_execution_sandbox.py
+++ b/tests/integration_test_tool_execution_sandbox.py
@@ -163,6 +163,26 @@ def get_env_tool(test_user):
     yield tool
 
 
+@pytest.fixture
+def get_warning_tool(test_user):
+    def warn_hello_world() -> str:
+        """
+        Simple function that warns hello world.
+
+        Returns:
+            str: hello world
+        """
+        import warnings
+
+        msg = "Hello World"
+        warnings.warn(msg)
+        return msg
+
+    tool = create_tool_from_func(warn_hello_world)
+    tool = ToolManager().create_or_update_tool(tool, test_user)
+    yield tool
+
+
 @pytest.fixture
 def list_tool(test_user):
     def create_list():
@@ -349,6 +369,22 @@ def test_local_sandbox_external_codebase(mock_e2b_api_key_none, external_codebas
     assert "Hello World" in result.stdout[0]
 
 
+@pytest.mark.local_sandbox
+def test_local_sandbox_with_venv_and_warnings_does_not_error(mock_e2b_api_key_none, get_warning_tool, test_user):
+    # Make the external codebase the sandbox config
+    manager = SandboxConfigManager(tool_settings)
+
+    # Set the sandbox to be within the external codebase path and use a venv
+    external_codebase_path = str(Path(__file__).parent / "test_tool_sandbox" / "restaurant_management_system")
+    local_sandbox_config = LocalSandboxConfig(sandbox_dir=external_codebase_path, use_venv=True)
+    config_create = SandboxConfigCreate(config=local_sandbox_config.model_dump())
+    manager.create_or_update_sandbox_config(sandbox_config_create=config_create, actor=test_user)
+
+    sandbox = ToolExecutionSandbox(get_warning_tool.name, {}, user_id=test_user.id)
+    result = sandbox.run()
+    assert result.func_return == "Hello World"
+
+
 # E2B sandbox tests