363 lines
18 KiB
Python
363 lines
18 KiB
Python
from typing import Dict, List, Optional
|
|
|
|
from letta.constants import LETTA_TOOL_EXECUTION_DIR
|
|
from letta.log import get_logger
|
|
from letta.orm.errors import NoResultFound
|
|
from letta.orm.sandbox_config import SandboxConfig as SandboxConfigModel, SandboxEnvironmentVariable as SandboxEnvVarModel
|
|
from letta.otel.tracing import trace_method
|
|
from letta.schemas.enums import PrimitiveType, SandboxType
|
|
from letta.schemas.environment_variables import (
|
|
SandboxEnvironmentVariable as PydanticEnvVar,
|
|
SandboxEnvironmentVariableCreate,
|
|
SandboxEnvironmentVariableUpdate,
|
|
)
|
|
from letta.schemas.sandbox_config import (
|
|
E2BSandboxConfig,
|
|
LocalSandboxConfig,
|
|
ModalSandboxConfig,
|
|
SandboxConfig as PydanticSandboxConfig,
|
|
SandboxConfigCreate,
|
|
SandboxConfigUpdate,
|
|
)
|
|
from letta.schemas.user import User as PydanticUser
|
|
from letta.server.db import db_registry
|
|
from letta.utils import enforce_types, printd
|
|
from letta.validators import raise_on_invalid_id
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class SandboxConfigManager:
|
|
"""Manager class to handle business logic related to SandboxConfig and SandboxEnvironmentVariable."""
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
def get_or_create_default_sandbox_config(self, sandbox_type: SandboxType, actor: PydanticUser) -> PydanticSandboxConfig:
|
|
sandbox_config = self.get_sandbox_config_by_type(sandbox_type, actor=actor)
|
|
if not sandbox_config:
|
|
logger.debug(f"Creating new sandbox config of type {sandbox_type}, none found for organization {actor.organization_id}.")
|
|
|
|
# Create the appropriate config type based on sandbox_type
|
|
# Using the actual model classes ensures Pydantic's Union type resolution works correctly
|
|
if sandbox_type == SandboxType.E2B:
|
|
default_config = E2BSandboxConfig()
|
|
elif sandbox_type == SandboxType.MODAL:
|
|
default_config = ModalSandboxConfig()
|
|
else:
|
|
# LOCAL sandbox type
|
|
default_local_sandbox_path = LETTA_TOOL_EXECUTION_DIR
|
|
default_config = LocalSandboxConfig(sandbox_dir=default_local_sandbox_path)
|
|
|
|
sandbox_config = self.create_or_update_sandbox_config(SandboxConfigCreate(config=default_config), actor=actor)
|
|
return sandbox_config
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def get_or_create_default_sandbox_config_async(self, sandbox_type: SandboxType, actor: PydanticUser) -> PydanticSandboxConfig:
|
|
sandbox_config = await self.get_sandbox_config_by_type_async(sandbox_type, actor=actor)
|
|
if not sandbox_config:
|
|
logger.debug(f"Creating new sandbox config of type {sandbox_type}, none found for organization {actor.organization_id}.")
|
|
|
|
# Create the appropriate config type based on sandbox_type
|
|
# Using the actual model classes ensures Pydantic's Union type resolution works correctly
|
|
if sandbox_type == SandboxType.E2B:
|
|
default_config = E2BSandboxConfig()
|
|
elif sandbox_type == SandboxType.MODAL:
|
|
default_config = ModalSandboxConfig()
|
|
else:
|
|
# LOCAL sandbox type
|
|
default_local_sandbox_path = LETTA_TOOL_EXECUTION_DIR
|
|
default_config = LocalSandboxConfig(sandbox_dir=default_local_sandbox_path)
|
|
|
|
sandbox_config = await self.create_or_update_sandbox_config_async(SandboxConfigCreate(config=default_config), actor=actor)
|
|
return sandbox_config
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def create_or_update_sandbox_config_async(
|
|
self, sandbox_config_create: SandboxConfigCreate, actor: PydanticUser
|
|
) -> PydanticSandboxConfig:
|
|
"""Create or update a sandbox configuration based on the PydanticSandboxConfig schema."""
|
|
config = sandbox_config_create.config
|
|
sandbox_type = config.type
|
|
sandbox_config = PydanticSandboxConfig(
|
|
type=sandbox_type, config=config.model_dump(exclude_none=True), organization_id=actor.organization_id
|
|
)
|
|
|
|
# Attempt to retrieve the existing sandbox configuration by type within the organization
|
|
db_sandbox = await self.get_sandbox_config_by_type_async(sandbox_config.type, actor=actor)
|
|
if db_sandbox:
|
|
# Prepare the update data, excluding fields that should not be reset
|
|
update_data = sandbox_config.model_dump(exclude_unset=True, exclude_none=True)
|
|
update_data = {key: value for key, value in update_data.items() if getattr(db_sandbox, key) != value}
|
|
|
|
# If there are changes, update the sandbox configuration
|
|
if update_data:
|
|
db_sandbox = await self.update_sandbox_config_async(db_sandbox.id, SandboxConfigUpdate(**update_data), actor)
|
|
else:
|
|
printd(
|
|
f"`create_or_update_sandbox_config` was called with user_id={actor.id}, organization_id={actor.organization_id}, "
|
|
f"type={sandbox_config.type}, but found existing configuration with nothing to update."
|
|
)
|
|
|
|
return db_sandbox
|
|
else:
|
|
# If the sandbox configuration doesn't exist, create a new one
|
|
async with db_registry.async_session() as session:
|
|
db_sandbox = SandboxConfigModel(**sandbox_config.model_dump(exclude_none=True))
|
|
await db_sandbox.create_async(session, actor=actor)
|
|
return db_sandbox.to_pydantic()
|
|
|
|
@enforce_types
|
|
@raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
|
|
@trace_method
|
|
async def update_sandbox_config_async(
|
|
self, sandbox_config_id: str, sandbox_update: SandboxConfigUpdate, actor: PydanticUser
|
|
) -> PydanticSandboxConfig:
|
|
"""Update an existing sandbox configuration."""
|
|
async with db_registry.async_session() as session:
|
|
sandbox = await SandboxConfigModel.read_async(db_session=session, identifier=sandbox_config_id, actor=actor)
|
|
# We need to check that the sandbox_update provided is the same type as the original sandbox
|
|
if sandbox.type != sandbox_update.config.type:
|
|
raise ValueError(
|
|
f"Mismatched type for sandbox config update: tried to update sandbox_config of type {sandbox.type} with config of type {sandbox_update.config.type}"
|
|
)
|
|
|
|
update_data = sandbox_update.model_dump(exclude_unset=True, exclude_none=True)
|
|
update_data = {key: value for key, value in update_data.items() if getattr(sandbox, key) != value}
|
|
|
|
if update_data:
|
|
for key, value in update_data.items():
|
|
setattr(sandbox, key, value)
|
|
await sandbox.update_async(db_session=session, actor=actor)
|
|
else:
|
|
printd(
|
|
f"`update_sandbox_config` called with user_id={actor.id}, organization_id={actor.organization_id}, "
|
|
f"name={sandbox.type}, but nothing to update."
|
|
)
|
|
return sandbox.to_pydantic()
|
|
|
|
@enforce_types
|
|
@raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
|
|
@trace_method
|
|
async def delete_sandbox_config_async(self, sandbox_config_id: str, actor: PydanticUser) -> PydanticSandboxConfig:
|
|
"""Delete a sandbox configuration by its ID."""
|
|
async with db_registry.async_session() as session:
|
|
sandbox = await SandboxConfigModel.read_async(db_session=session, identifier=sandbox_config_id, actor=actor)
|
|
await sandbox.hard_delete_async(db_session=session, actor=actor)
|
|
return sandbox.to_pydantic()
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def list_sandbox_configs_async(
|
|
self,
|
|
actor: PydanticUser,
|
|
after: Optional[str] = None,
|
|
limit: Optional[int] = 50,
|
|
sandbox_type: Optional[SandboxType] = None,
|
|
) -> List[PydanticSandboxConfig]:
|
|
"""List all sandbox configurations with optional pagination."""
|
|
kwargs = {"organization_id": actor.organization_id}
|
|
if sandbox_type:
|
|
kwargs.update({"type": sandbox_type})
|
|
|
|
async with db_registry.async_session() as session:
|
|
sandboxes = await SandboxConfigModel.list_async(db_session=session, after=after, limit=limit, **kwargs)
|
|
return [sandbox.to_pydantic() for sandbox in sandboxes]
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def get_sandbox_config_by_type_async(self, type: SandboxType, actor: PydanticUser) -> Optional[PydanticSandboxConfig]:
|
|
"""Retrieve a sandbox config by its type."""
|
|
async with db_registry.async_session() as session:
|
|
try:
|
|
sandboxes = await SandboxConfigModel.list_async(
|
|
db_session=session,
|
|
type=type,
|
|
organization_id=actor.organization_id,
|
|
limit=1,
|
|
)
|
|
if sandboxes:
|
|
return sandboxes[0].to_pydantic()
|
|
return None
|
|
except NoResultFound:
|
|
return None
|
|
|
|
@enforce_types
|
|
@raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
|
|
@trace_method
|
|
async def create_sandbox_env_var_async(
|
|
self, env_var_create: SandboxEnvironmentVariableCreate, sandbox_config_id: str, actor: PydanticUser
|
|
) -> PydanticEnvVar:
|
|
"""Create a new sandbox environment variable."""
|
|
env_var = PydanticEnvVar(**env_var_create.model_dump(), sandbox_config_id=sandbox_config_id, organization_id=actor.organization_id)
|
|
|
|
db_env_var = await self.get_sandbox_env_var_by_key_and_sandbox_config_id_async(env_var.key, env_var.sandbox_config_id, actor=actor)
|
|
if db_env_var:
|
|
update_data = env_var.model_dump(exclude_unset=True, exclude_none=True)
|
|
update_data = {key: value for key, value in update_data.items() if getattr(db_env_var, key) != value}
|
|
# If there are changes, update the environment variable
|
|
if update_data:
|
|
db_env_var = await self.update_sandbox_env_var_async(db_env_var.id, SandboxEnvironmentVariableUpdate(**update_data), actor)
|
|
else:
|
|
printd(
|
|
f"`create_or_update_sandbox_env_var` was called with user_id={actor.id}, organization_id={actor.organization_id}, "
|
|
f"key={env_var.key}, but found existing variable with nothing to update."
|
|
)
|
|
|
|
return db_env_var
|
|
else:
|
|
async with db_registry.async_session() as session:
|
|
# Encrypt the value before storing (async to avoid blocking event loop)
|
|
from letta.schemas.secret import Secret
|
|
|
|
if env_var.value:
|
|
env_var.value_enc = await Secret.from_plaintext_async(env_var.value)
|
|
env_var.value = "" # Don't store plaintext, use empty string for NOT NULL constraint
|
|
|
|
env_var = SandboxEnvVarModel(**env_var.model_dump(to_orm=True))
|
|
await env_var.create_async(session, actor=actor)
|
|
return await PydanticEnvVar.from_orm_async(env_var)
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def update_sandbox_env_var_async(
|
|
self, env_var_id: str, env_var_update: SandboxEnvironmentVariableUpdate, actor: PydanticUser
|
|
) -> PydanticEnvVar:
|
|
"""Update an existing sandbox environment variable."""
|
|
async with db_registry.async_session() as session:
|
|
env_var = await SandboxEnvVarModel.read_async(db_session=session, identifier=env_var_id, actor=actor)
|
|
update_data = env_var_update.model_dump(to_orm=True, exclude_unset=True, exclude_none=True)
|
|
|
|
# Handle encryption for value if provided
|
|
# Only re-encrypt if the value has actually changed
|
|
if "value" in update_data and update_data["value"] is not None:
|
|
from letta.schemas.secret import Secret
|
|
|
|
# Check if value changed by comparing with existing encrypted value
|
|
existing_value = None
|
|
if env_var.value_enc:
|
|
existing_secret = Secret.from_encrypted(env_var.value_enc)
|
|
existing_value = await existing_secret.get_plaintext_async()
|
|
|
|
# Only re-encrypt if different (async to avoid blocking event loop)
|
|
if existing_value != update_data["value"]:
|
|
value_secret = await Secret.from_plaintext_async(update_data["value"])
|
|
env_var.value_enc = value_secret.get_encrypted()
|
|
# Don't store plaintext anymore
|
|
|
|
# Remove from update_data since we set directly on env_var
|
|
update_data.pop("value", None)
|
|
update_data.pop("value_enc", None)
|
|
|
|
# Apply remaining updates
|
|
update_data = {key: value for key, value in update_data.items() if getattr(env_var, key) != value}
|
|
|
|
if update_data:
|
|
for key, value in update_data.items():
|
|
setattr(env_var, key, value)
|
|
await env_var.update_async(db_session=session, actor=actor)
|
|
else:
|
|
printd(
|
|
f"`update_sandbox_env_var` called with user_id={actor.id}, organization_id={actor.organization_id}, "
|
|
f"key={env_var.key}, but nothing to update."
|
|
)
|
|
return await PydanticEnvVar.from_orm_async(env_var)
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def delete_sandbox_env_var_async(self, env_var_id: str, actor: PydanticUser) -> PydanticEnvVar:
|
|
"""Delete a sandbox environment variable by its ID."""
|
|
async with db_registry.async_session() as session:
|
|
env_var = await SandboxEnvVarModel.read_async(db_session=session, identifier=env_var_id, actor=actor)
|
|
await env_var.hard_delete_async(db_session=session, actor=actor)
|
|
return await PydanticEnvVar.from_orm_async(env_var)
|
|
|
|
@enforce_types
|
|
@raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
|
|
@trace_method
|
|
async def list_sandbox_env_vars_async(
|
|
self,
|
|
sandbox_config_id: str,
|
|
actor: PydanticUser,
|
|
after: Optional[str] = None,
|
|
limit: Optional[int] = 50,
|
|
) -> List[PydanticEnvVar]:
|
|
"""List all sandbox environment variables with optional pagination."""
|
|
async with db_registry.async_session() as session:
|
|
env_vars = await SandboxEnvVarModel.list_async(
|
|
db_session=session,
|
|
after=after,
|
|
limit=limit,
|
|
organization_id=actor.organization_id,
|
|
sandbox_config_id=sandbox_config_id,
|
|
)
|
|
result = []
|
|
for env_var in env_vars:
|
|
result.append(await PydanticEnvVar.from_orm_async(env_var))
|
|
return result
|
|
|
|
@enforce_types
|
|
@trace_method
|
|
async def list_sandbox_env_vars_by_key_async(
|
|
self, key: str, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50
|
|
) -> List[PydanticEnvVar]:
|
|
"""List all sandbox environment variables with optional pagination."""
|
|
async with db_registry.async_session() as session:
|
|
env_vars = await SandboxEnvVarModel.list_async(
|
|
db_session=session,
|
|
after=after,
|
|
limit=limit,
|
|
organization_id=actor.organization_id,
|
|
key=key,
|
|
)
|
|
result = []
|
|
for env_var in env_vars:
|
|
result.append(await PydanticEnvVar.from_orm_async(env_var))
|
|
return result
|
|
|
|
@enforce_types
|
|
@raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
|
|
@trace_method
|
|
def get_sandbox_env_vars_as_dict(
|
|
self, sandbox_config_id: str, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50
|
|
) -> Dict[str, str]:
|
|
env_vars = self.list_sandbox_env_vars(sandbox_config_id, actor, after, limit)
|
|
result = {}
|
|
for env_var in env_vars:
|
|
# Decrypt the value before returning
|
|
result[env_var.key] = env_var.value_enc.get_plaintext() if env_var.value_enc else None
|
|
return result
|
|
|
|
@enforce_types
|
|
@raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
|
|
@trace_method
|
|
async def get_sandbox_env_vars_as_dict_async(
|
|
self, sandbox_config_id: str, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50
|
|
) -> Dict[str, str]:
|
|
env_vars = await self.list_sandbox_env_vars_async(sandbox_config_id, actor, after, limit)
|
|
# Values are already decrypted via from_orm_async
|
|
return {env_var.key: env_var.value for env_var in env_vars}
|
|
|
|
@enforce_types
|
|
@raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
|
|
@trace_method
|
|
async def get_sandbox_env_var_by_key_and_sandbox_config_id_async(
|
|
self, key: str, sandbox_config_id: str, actor: PydanticUser
|
|
) -> Optional[PydanticEnvVar]:
|
|
"""Retrieve a sandbox environment variable by its key and sandbox_config_id."""
|
|
async with db_registry.async_session() as session:
|
|
try:
|
|
env_var = await SandboxEnvVarModel.list_async(
|
|
db_session=session,
|
|
key=key,
|
|
sandbox_config_id=sandbox_config_id,
|
|
organization_id=actor.organization_id,
|
|
limit=1,
|
|
)
|
|
if env_var:
|
|
return await PydanticEnvVar.from_orm_async(env_var[0])
|
|
return None
|
|
except NoResultFound:
|
|
return None
|