Fimeg/letta-server
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

feat: Create local sandbox config from env vars for OSS users (#1910)

Browse Source
This commit is contained in:
Matthew Zhou
2025-04-28 13:01:38 -07:00
committed by GitHub
parent 6609372676
commit 59aefc322a
11 changed files with 111 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
letta/schemas/sandbox_config.py
View File

@@ -47,14 +47,14 @@ class PipRequirement(BaseModel):
class LocalSandboxConfig(BaseModel):
sandbox_dir: Optional[str] = Field(None, description="Directory for the sandbox environment.")
force_create_venv: bool = Field(False, description="Whether or not to use the venv, or run directly in the same run loop.")
use_venv: bool = Field(False, description="Whether or not to use the venv, or run directly in the same run loop.")
venv_name: str = Field(
"venv",
description="The name for the venv in the sandbox directory. We first search for an existing venv with this name, otherwise, we make it from the requirements.txt.",
)
pip_requirements: List[PipRequirement] = Field(
default_factory=list,
description="List of pip packages to install with mandatory name and optional version following semantic versioning. This only is considered when force_create_venv is True.",
description="List of pip packages to install with mandatory name and optional version following semantic versioning. This only is considered when use_venv is True.",
)
@property
@@ -69,8 +69,8 @@ class LocalSandboxConfig(BaseModel):
return data
if data.get("sandbox_dir") is None:
if tool_settings.local_sandbox_dir:
data["sandbox_dir"] = tool_settings.local_sandbox_dir
if tool_settings.tool_exec_dir:
data["sandbox_dir"] = tool_settings.tool_exec_dir
else:
data["sandbox_dir"] = LETTA_TOOL_EXECUTION_DIR

27
letta/server/server.py
View File

@@ -6,6 +6,7 @@ import traceback
import warnings
from abc import abstractmethod
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
from anthropic import AsyncAnthropic
@@ -19,6 +20,7 @@ import letta.server.utils as server_utils
import letta.system as system
from letta.agent import Agent, save_agent
from letta.config import LettaConfig
from letta.constants import LETTA_TOOL_EXECUTION_DIR
from letta.data_sources.connectors import DataConnector, load_data
from letta.errors import HandleNotFoundError
from letta.functions.mcp_client.base_client import BaseMCPClient
@@ -70,7 +72,7 @@ from letta.schemas.providers import (
VLLMCompletionsProvider,
XAIProvider,
)
from letta.schemas.sandbox_config import SandboxType
from letta.schemas.sandbox_config import LocalSandboxConfig, SandboxConfigCreate, SandboxType
from letta.schemas.source import Source
from letta.schemas.tool import Tool
from letta.schemas.usage import LettaUsageStatistics
@@ -229,6 +231,29 @@ class SyncServer(Server):
actor=self.default_user,
)
# For OSS users, create a local sandbox config
oss_default_user = self.user_manager.get_default_user()
use_venv = False if not tool_settings.tool_exec_venv_name else True
venv_name = tool_settings.tool_exec_venv_name or "venv"
tool_dir = tool_settings.tool_exec_dir or LETTA_TOOL_EXECUTION_DIR
venv_dir = Path(tool_dir) / venv_name
if not Path(tool_dir).is_dir():
logger.error(f"Provided LETTA_TOOL_SANDBOX_DIR is not a valid directory: {tool_dir}")
else:
if tool_settings.tool_exec_venv_name and not venv_dir.is_dir():
logger.warning(
f"Provided LETTA_TOOL_SANDBOX_VENV_NAME is not a valid venv ({venv_dir}), one will be created for you during tool execution."
)
sandbox_config_create = SandboxConfigCreate(
config=LocalSandboxConfig(sandbox_dir=tool_settings.tool_exec_dir, use_venv=use_venv, venv_name=venv_name)
)
sandbox_config = self.sandbox_config_manager.create_or_update_sandbox_config(
sandbox_config_create=sandbox_config_create, actor=oss_default_user
)
logger.info(f"Successfully created default local sandbox config:\n{sandbox_config.get_local_config().model_dump()}")
# collect providers (always has Letta as a default)
self._enabled_providers: List[Provider] = [LettaProvider()]
if model_settings.openai_api_key:

6
letta/services/helpers/tool_execution_helper.py
View File

@@ -24,7 +24,7 @@ def find_python_executable(local_configs: LocalSandboxConfig) -> str:
"""
sandbox_dir = os.path.expanduser(local_configs.sandbox_dir) # Expand tilde
if not local_configs.force_create_venv:
if not local_configs.use_venv:
return "python.exe" if platform.system().lower().startswith("win") else "python3"
venv_path = os.path.join(sandbox_dir, local_configs.venv_name)
@@ -96,7 +96,7 @@ def install_pip_requirements_for_sandbox(
python_exec = find_python_executable(local_configs)
# If using a virtual environment, upgrade pip before installing dependencies.
if local_configs.force_create_venv:
if local_configs.use_venv:
ensure_pip_is_up_to_date(python_exec, env=env)
# Construct package list
@@ -108,7 +108,7 @@ def install_pip_requirements_for_sandbox(
pip_cmd.append("--upgrade")
pip_cmd += packages
if user_install_if_no_venv and not local_configs.force_create_venv:
if user_install_if_no_venv and not local_configs.use_venv:
pip_cmd.append("--user")
run_subprocess(pip_cmd, env=env, fail_msg=f"Failed to install packages: {', '.join(packages)}")

22
letta/services/tool_executor/tool_execution_sandbox.py
View File

@@ -144,7 +144,7 @@ class ToolExecutionSandbox:
# Write the code to a temp file in the sandbox_dir
with tempfile.NamedTemporaryFile(mode="w", dir=local_configs.sandbox_dir, suffix=".py", delete=False) as temp_file:
if local_configs.force_create_venv:
if local_configs.use_venv:
# If using venv, we need to wrap with special string markers to separate out the output and the stdout (since it is all in stdout)
code = self.generate_execution_script(agent_state=agent_state, wrap_print_with_markers=True)
else:
@@ -154,7 +154,7 @@ class ToolExecutionSandbox:
temp_file.flush()
temp_file_path = temp_file.name
try:
if local_configs.force_create_venv:
if local_configs.use_venv:
return self.run_local_dir_sandbox_venv(sbx_config, env, temp_file_path)
else:
return self.run_local_dir_sandbox_directly(sbx_config, env, temp_file_path)
@@ -220,7 +220,11 @@ class ToolExecutionSandbox:
)
except subprocess.CalledProcessError as e:
with open(temp_file_path, "r") as f:
code = f.read()
logger.error(f"Executing tool {self.tool_name} has process error: {e}")
logger.error(f"Logging out tool {self.tool_name} auto-generated code for debugging: \n\n{code}")
func_return = get_friendly_error_msg(
function_name=self.tool_name,
exception_name=type(e).__name__,
@@ -447,6 +451,11 @@ class ToolExecutionSandbox:
Returns:
code (str): The generated code strong
"""
if "agent_state" in self.parse_function_arguments(self.tool.source_code, self.tool.name):
inject_agent_state = True
else:
inject_agent_state = False
# dump JSON representation of agent state to re-load
code = "from typing import *\n"
code += "import pickle\n"
@@ -454,7 +463,7 @@ class ToolExecutionSandbox:
code += "import base64\n"
# imports to support agent state
if agent_state:
if inject_agent_state:
code += "import letta\n"
code += "from letta import * \n"
import pickle
@@ -467,7 +476,7 @@ class ToolExecutionSandbox:
code += schema_code + "\n"
# load the agent state
if agent_state:
if inject_agent_state:
agent_state_pickle = pickle.dumps(agent_state)
code += f"agent_state = pickle.loads({agent_state_pickle})\n"
else:
@@ -483,11 +492,6 @@ class ToolExecutionSandbox:
for param in self.args:
code += self.initialize_param(param, self.args[param])
if "agent_state" in self.parse_function_arguments(self.tool.source_code, self.tool.name):
inject_agent_state = True
else:
inject_agent_state = False
code += "\n" + self.tool.source_code + "\n"
# TODO: handle wrapped print

8
letta/services/tool_sandbox/local_sandbox.py
View File

@@ -69,7 +69,7 @@ class AsyncToolSandboxLocal(AsyncToolSandboxBase):
else:
sbx_config = self.sandbox_config_manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.LOCAL, actor=self.user)
local_configs = sbx_config.get_local_config()
force_create_venv = local_configs.force_create_venv
use_venv = local_configs.use_venv
# Prepare environment variables
env = os.environ.copy()
@@ -92,7 +92,7 @@ class AsyncToolSandboxLocal(AsyncToolSandboxBase):
# If using a virtual environment, ensure it's prepared in parallel
venv_preparation_task = None
if force_create_venv:
if use_venv:
venv_path = str(os.path.join(sandbox_dir, local_configs.venv_name))
venv_preparation_task = asyncio.create_task(self._prepare_venv(local_configs, venv_path, env))
@@ -110,7 +110,7 @@ class AsyncToolSandboxLocal(AsyncToolSandboxBase):
# Determine the python executable and environment for the subprocess
exec_env = env.copy()
if force_create_venv:
if use_venv:
venv_path = str(os.path.join(sandbox_dir, local_configs.venv_name))
python_executable = find_python_executable(local_configs)
exec_env["VIRTUAL_ENV"] = venv_path
@@ -174,7 +174,7 @@ class AsyncToolSandboxLocal(AsyncToolSandboxBase):
)
try:
stdout_bytes, stderr_bytes = await asyncio.wait_for(process.communicate(), timeout=tool_settings.local_sandbox_timeout)
stdout_bytes, stderr_bytes = await asyncio.wait_for(process.communicate(), timeout=tool_settings.tool_sandbox_timeout)
except asyncio.TimeoutError:
# Terminate the process on timeout
if process.returncode is None:

7
letta/services/user_manager.py
View File

@@ -84,8 +84,11 @@ class UserManager:
@enforce_types
def get_default_user(self) -> PydanticUser:
"""Fetch the default user."""
return self.get_user_by_id(self.DEFAULT_USER_ID)
"""Fetch the default user. If it doesn't exist, create it."""
try:
return self.get_user_by_id(self.DEFAULT_USER_ID)
except NoResultFound:
return self.create_default_user()
@enforce_types
def get_user_or_default(self, user_id: Optional[str] = None):

5
letta/settings.py
View File

@@ -16,8 +16,9 @@ class ToolSettings(BaseSettings):
e2b_sandbox_template_id: Optional[str] = None # Updated manually
# Local Sandbox configurations
local_sandbox_dir: Optional[str] = None
local_sandbox_timeout: float = 180
tool_exec_dir: Optional[str] = None
tool_sandbox_timeout: float = 180
tool_exec_venv_name: Optional[str] = None
# MCP settings
mcp_connect_to_server_timeout: float = 30.0

8
tests/integration_test_async_tool_sandbox.py
View File

@@ -221,7 +221,7 @@ def custom_test_sandbox_config(test_user):
external_codebase_path = str(Path(__file__).parent / "test_tool_sandbox" / "restaurant_management_system")
# tqdm is used in this codebase, but NOT in the requirements.txt, this tests that we can successfully install pip requirements
local_sandbox_config = LocalSandboxConfig(
sandbox_dir=external_codebase_path, force_create_venv=True, pip_requirements=[PipRequirement(name="tqdm")]
sandbox_dir=external_codebase_path, use_venv=True, pip_requirements=[PipRequirement(name="tqdm")]
)
# Create the sandbox configuration
@@ -366,7 +366,7 @@ async def test_local_sandbox_with_venv_errors(disable_e2b_api_key, custom_test_s
async def test_local_sandbox_with_venv_pip_installs_basic(disable_e2b_api_key, cowsay_tool, test_user):
manager = SandboxConfigManager()
config_create = SandboxConfigCreate(
config=LocalSandboxConfig(force_create_venv=True, pip_requirements=[PipRequirement(name="cowsay")]).model_dump()
config=LocalSandboxConfig(use_venv=True, pip_requirements=[PipRequirement(name="cowsay")]).model_dump()
)
config = manager.create_or_update_sandbox_config(config_create, test_user)
@@ -385,7 +385,7 @@ async def test_local_sandbox_with_venv_pip_installs_basic(disable_e2b_api_key, c
@pytest.mark.e2b_sandbox
async def test_local_sandbox_with_venv_pip_installs_with_update(disable_e2b_api_key, cowsay_tool, test_user):
manager = SandboxConfigManager()
config_create = SandboxConfigCreate(config=LocalSandboxConfig(force_create_venv=True).model_dump())
config_create = SandboxConfigCreate(config=LocalSandboxConfig(use_venv=True).model_dump())
config = manager.create_or_update_sandbox_config(config_create, test_user)
key = "secret_word"
@@ -400,7 +400,7 @@ async def test_local_sandbox_with_venv_pip_installs_with_update(disable_e2b_api_
assert "No module named 'cowsay'" in result.stderr[0]
config_create = SandboxConfigCreate(
config=LocalSandboxConfig(force_create_venv=True, pip_requirements=[PipRequirement(name="cowsay")]).model_dump()
config=LocalSandboxConfig(use_venv=True, pip_requirements=[PipRequirement(name="cowsay")]).model_dump()
)
manager.create_or_update_sandbox_config(config_create, test_user)

8
tests/integration_test_tool_execution_sandbox.py
View File

@@ -220,7 +220,7 @@ def custom_test_sandbox_config(test_user):
external_codebase_path = str(Path(__file__).parent / "test_tool_sandbox" / "restaurant_management_system")
# tqdm is used in this codebase, but NOT in the requirements.txt, this tests that we can successfully install pip requirements
local_sandbox_config = LocalSandboxConfig(
sandbox_dir=external_codebase_path, force_create_venv=True, pip_requirements=[PipRequirement(name="tqdm")]
sandbox_dir=external_codebase_path, use_venv=True, pip_requirements=[PipRequirement(name="tqdm")]
)
# Create the sandbox configuration
@@ -382,7 +382,7 @@ def test_local_sandbox_with_venv_errors(disable_e2b_api_key, custom_test_sandbox
def test_local_sandbox_with_venv_pip_installs_basic(disable_e2b_api_key, cowsay_tool, test_user):
manager = SandboxConfigManager()
config_create = SandboxConfigCreate(
config=LocalSandboxConfig(force_create_venv=True, pip_requirements=[PipRequirement(name="cowsay")]).model_dump()
config=LocalSandboxConfig(use_venv=True, pip_requirements=[PipRequirement(name="cowsay")]).model_dump()
)
config = manager.create_or_update_sandbox_config(config_create, test_user)
@@ -401,7 +401,7 @@ def test_local_sandbox_with_venv_pip_installs_basic(disable_e2b_api_key, cowsay_
@pytest.mark.e2b_sandbox
def test_local_sandbox_with_venv_pip_installs_with_update(disable_e2b_api_key, cowsay_tool, test_user):
manager = SandboxConfigManager()
config_create = SandboxConfigCreate(config=LocalSandboxConfig(force_create_venv=True).model_dump())
config_create = SandboxConfigCreate(config=LocalSandboxConfig(use_venv=True).model_dump())
config = manager.create_or_update_sandbox_config(config_create, test_user)
# Add an environment variable
@@ -421,7 +421,7 @@ def test_local_sandbox_with_venv_pip_installs_with_update(disable_e2b_api_key, c
# Now update the SandboxConfig
config_create = SandboxConfigCreate(
config=LocalSandboxConfig(force_create_venv=True, pip_requirements=[PipRequirement(name="cowsay")]).model_dump()
config=LocalSandboxConfig(use_venv=True, pip_requirements=[PipRequirement(name="cowsay")]).model_dump()
)
manager.create_or_update_sandbox_config(config_create, test_user)

2
tests/test_managers.py
View File

@@ -3852,7 +3852,7 @@ def test_create_local_sandbox_config_defaults(server: SyncServer, default_user):
# Assertions
assert created_config.type == SandboxType.LOCAL
assert created_config.get_local_config() == sandbox_config_create.config
assert created_config.get_local_config().sandbox_dir in {LETTA_TOOL_EXECUTION_DIR, tool_settings.local_sandbox_dir}
assert created_config.get_local_config().sandbox_dir in {LETTA_TOOL_EXECUTION_DIR, tool_settings.tool_exec_dir}
assert created_config.organization_id == default_user.organization_id

45
tests/test_server.py
View File

@@ -4,18 +4,20 @@ import shutil
import uuid
import warnings
from typing import List, Tuple
from unittest.mock import patch
import pytest
from sqlalchemy import delete
import letta.utils as utils
from letta.constants import BASE_MEMORY_TOOLS, BASE_TOOLS
from letta.constants import BASE_MEMORY_TOOLS, BASE_TOOLS, LETTA_DIR, LETTA_TOOL_EXECUTION_DIR
from letta.orm import Provider, Step
from letta.schemas.block import CreateBlock
from letta.schemas.enums import MessageRole
from letta.schemas.letta_message import LettaMessage, ReasoningMessage, SystemMessage, ToolCallMessage, ToolReturnMessage, UserMessage
from letta.schemas.llm_config import LLMConfig
from letta.schemas.providers import Provider as PydanticProvider
from letta.schemas.sandbox_config import SandboxType
from letta.schemas.user import User
utils.DEBUG = True
@@ -1299,3 +1301,44 @@ def test_unique_handles_for_provider_configs(server: SyncServer):
embeddings = server.list_embedding_models()
embedding_handles = [embedding.handle for embedding in embeddings]
assert sorted(embedding_handles) == sorted(list(set(embedding_handles))), "All embeddings should have unique handles"
def test_make_default_local_sandbox_config():
venv_name = "test"
default_venv_name = "venv"
# --- Case 1: tool_exec_dir and tool_exec_venv_name are both explicitly set ---
with patch("letta.settings.tool_settings.tool_exec_dir", LETTA_DIR):
with patch("letta.settings.tool_settings.tool_exec_venv_name", venv_name):
server = SyncServer()
actor = server.user_manager.get_default_user()
local_config = server.sandbox_config_manager.get_or_create_default_sandbox_config(
sandbox_type=SandboxType.LOCAL, actor=actor
).get_local_config()
assert local_config.sandbox_dir == LETTA_DIR
assert local_config.venv_name == venv_name
assert local_config.use_venv == True
# --- Case 2: only tool_exec_dir is set (no custom venv_name provided) ---
with patch("letta.settings.tool_settings.tool_exec_dir", LETTA_DIR):
server = SyncServer()
actor = server.user_manager.get_default_user()
local_config = server.sandbox_config_manager.get_or_create_default_sandbox_config(
sandbox_type=SandboxType.LOCAL, actor=actor
).get_local_config()
assert local_config.sandbox_dir == LETTA_DIR
assert local_config.venv_name == default_venv_name # falls back to default
assert local_config.use_venv == False # no custom venv name, so no venv usage
# --- Case 3: neither tool_exec_dir nor tool_exec_venv_name is set (default fallback behavior) ---
server = SyncServer()
actor = server.user_manager.get_default_user()
local_config = server.sandbox_config_manager.get_or_create_default_sandbox_config(
sandbox_type=SandboxType.LOCAL, actor=actor
).get_local_config()
assert local_config.sandbox_dir == LETTA_TOOL_EXECUTION_DIR
assert local_config.venv_name == default_venv_name
assert local_config.use_venv == False