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

feat: Inject per agent environment variables into sandbox (#514)

Browse Source
This commit is contained in:
Matthew Zhou
2025-01-06 13:04:17 -10:00
committed by GitHub
parent 926a6fa07a
commit cf47802a8c
3 changed files with 90 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
letta/schemas/agent.py
View File

@@ -83,6 +83,13 @@ class AgentState(OrmMetadataBase, validate_assignment=True):
..., description="The environment variables for tool execution specific to this agent."
)
def get_agent_env_vars_as_dict(self) -> Dict[str, str]:
# Get environment variables for this agent specifically
per_agent_env_vars = {}
for agent_env_var_obj in self.tool_exec_environment_variables:
per_agent_env_vars[agent_env_var_obj.key] = agent_env_var_obj.value
return per_agent_env_vars
class CreateAgent(BaseModel, validate_assignment=True): #
# all optional as server can generate defaults

12
letta/services/tool_execution_sandbox.py
View File

@@ -98,7 +98,7 @@ class ToolExecutionSandbox:
os.environ.clear()
os.environ.update(original_env) # Restore original environment variables
def run_local_dir_sandbox(self, agent_state: AgentState) -> SandboxRunResult:
def run_local_dir_sandbox(self, agent_state: Optional[AgentState] = None) -> SandboxRunResult:
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()
@@ -107,6 +107,10 @@ class ToolExecutionSandbox:
env = os.environ.copy()
env.update(env_vars)
# Get environment variables for this agent specifically
if agent_state:
env.update(agent_state.get_agent_env_vars_as_dict())
# Safety checks
if not os.path.isdir(local_configs.sandbox_dir):
raise FileNotFoundError(f"Sandbox directory does not exist: {local_configs.sandbox_dir}")
@@ -273,7 +277,7 @@ class ToolExecutionSandbox:
# e2b sandbox specific functions
def run_e2b_sandbox(self, agent_state: AgentState) -> SandboxRunResult:
def run_e2b_sandbox(self, agent_state: Optional[AgentState] = None) -> SandboxRunResult:
sbx_config = self.sandbox_config_manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.E2B, actor=self.user)
sbx = self.get_running_e2b_sandbox_with_same_state(sbx_config)
if not sbx or self.force_recreate:
@@ -292,6 +296,10 @@ class ToolExecutionSandbox:
# Get environment variables for the sandbox
# TODO: We set limit to 100 here, but maybe we want it uncapped? Realistically this should be fine.
env_vars = self.sandbox_config_manager.get_sandbox_env_vars_as_dict(sandbox_config_id=sbx_config.id, actor=self.user, limit=100)
# Get environment variables for this agent specifically
if agent_state:
env_vars.update(agent_state.get_agent_env_vars_as_dict())
code = self.generate_execution_script(agent_state=agent_state)
execution = sbx.run_code(code, envs=env_vars)

75
tests/integration_test_tool_execution_sandbox.py
View File

@@ -12,7 +12,7 @@ from letta.functions.function_sets.base import core_memory_append, core_memory_r
from letta.orm.sandbox_config import SandboxConfig, SandboxEnvironmentVariable
from letta.schemas.agent import AgentState
from letta.schemas.embedding_config import EmbeddingConfig
from letta.schemas.environment_variables import SandboxEnvironmentVariableCreate
from letta.schemas.environment_variables import AgentEnvironmentVariable, SandboxEnvironmentVariableCreate
from letta.schemas.llm_config import LLMConfig
from letta.schemas.memory import ChatMemory
from letta.schemas.organization import Organization
@@ -323,6 +323,41 @@ def test_local_sandbox_env(mock_e2b_api_key_none, get_env_tool, test_user):
assert long_random_string in result.func_return
@pytest.mark.local_sandbox
def test_local_sandbox_per_agent_env(mock_e2b_api_key_none, get_env_tool, agent_state, test_user):
manager = SandboxConfigManager(tool_settings)
key = "secret_word"
# Make a custom local sandbox config
sandbox_dir = str(Path(__file__).parent / "test_tool_sandbox")
config_create = SandboxConfigCreate(config=LocalSandboxConfig(sandbox_dir=sandbox_dir).model_dump())
config = manager.create_or_update_sandbox_config(config_create, test_user)
# Make a environment variable with a long random string
# Note: This has an overlapping key with agent state's environment variables
# We expect that the agent's env var supersedes this
wrong_long_random_string = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(20))
manager.create_sandbox_env_var(
SandboxEnvironmentVariableCreate(key=key, value=wrong_long_random_string), sandbox_config_id=config.id, actor=test_user
)
# Make a environment variable with a long random string and put into agent state
correct_long_random_string = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(20))
agent_state.tool_exec_environment_variables = [
AgentEnvironmentVariable(key=key, value=correct_long_random_string, agent_id=agent_state.id)
]
# Create tool and args
args = {}
# Run the custom sandbox
sandbox = ToolExecutionSandbox(get_env_tool.name, args, user=test_user)
result = sandbox.run(agent_state=agent_state)
assert wrong_long_random_string not in result.func_return
assert correct_long_random_string in result.func_return
@pytest.mark.local_sandbox
def test_local_sandbox_e2e_composio_star_github(mock_e2b_api_key_none, check_composio_key_set, composio_github_star_tool, test_user):
# Add the composio key
@@ -470,6 +505,42 @@ def test_e2b_sandbox_inject_env_var_existing_sandbox(check_e2b_key_is_set, get_e
assert long_random_string in result.func_return
# TODO: There is a near dupe of this test above for local sandbox - we should try to make it parameterized tests to minimize code bloat
@pytest.mark.e2b_sandbox
def test_e2b_sandbox_per_agent_env(check_e2b_key_is_set, get_env_tool, agent_state, test_user):
manager = SandboxConfigManager(tool_settings)
key = "secret_word"
# Make a custom local sandbox config
sandbox_dir = str(Path(__file__).parent / "test_tool_sandbox")
config_create = SandboxConfigCreate(config=LocalSandboxConfig(sandbox_dir=sandbox_dir).model_dump())
config = manager.create_or_update_sandbox_config(config_create, test_user)
# Make a environment variable with a long random string
# Note: This has an overlapping key with agent state's environment variables
# We expect that the agent's env var supersedes this
wrong_long_random_string = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(20))
manager.create_sandbox_env_var(
SandboxEnvironmentVariableCreate(key=key, value=wrong_long_random_string), sandbox_config_id=config.id, actor=test_user
)
# Make a environment variable with a long random string and put into agent state
correct_long_random_string = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(20))
agent_state.tool_exec_environment_variables = [
AgentEnvironmentVariable(key=key, value=correct_long_random_string, agent_id=agent_state.id)
]
# Create tool and args
args = {}
# Run the custom sandbox
sandbox = ToolExecutionSandbox(get_env_tool.name, args, user=test_user)
result = sandbox.run(agent_state=agent_state)
assert wrong_long_random_string not in result.func_return
assert correct_long_random_string in result.func_return
@pytest.mark.e2b_sandbox
def test_e2b_sandbox_config_change_force_recreates_sandbox(check_e2b_key_is_set, list_tool, test_user):
manager = SandboxConfigManager(tool_settings)
@@ -506,7 +577,7 @@ def test_e2b_sandbox_with_list_rv(check_e2b_key_is_set, list_tool, test_user):
assert len(result.func_return) == 5
@pytest.mark.e2b_sandboxfunc
@pytest.mark.e2b_sandbox
def test_e2b_e2e_composio_star_github(check_e2b_key_is_set, check_composio_key_set, composio_github_star_tool, test_user):
# Add the composio key
manager = SandboxConfigManager(tool_settings)