diff --git a/letta/agent.py b/letta/agent.py index 52f7de37..f1d1f3d0 100644 --- a/letta/agent.py +++ b/letta/agent.py @@ -1628,7 +1628,7 @@ class Agent(BaseAgent): action_name = generate_composio_action_from_func_name(target_letta_tool.name) # Get entity ID from the agent_state entity_id = None - for env_var in self.agent_state.tool_exec_environment_variables: + for env_var in self.agent_state.secrets: if env_var.key == COMPOSIO_ENTITY_ENV_VAR_KEY: entity_id = env_var.value # Get composio_api_key diff --git a/letta/agents/letta_agent.py b/letta/agents/letta_agent.py index 2551e6c2..91034744 100644 --- a/letta/agents/letta_agent.py +++ b/letta/agents/letta_agent.py @@ -1874,7 +1874,7 @@ class LettaAgent(BaseAgent): start_time = get_utc_timestamp_ns() agent_step_span.add_event(name="tool_execution_started") - sandbox_env_vars = {var.key: var.value for var in agent_state.tool_exec_environment_variables} + sandbox_env_vars = {var.key: var.value for var in agent_state.secrets} tool_execution_manager = ToolExecutionManager( agent_state=agent_state, message_manager=self.message_manager, diff --git a/letta/agents/letta_agent_v2.py b/letta/agents/letta_agent_v2.py index 09a02be3..524f5d6b 100644 --- a/letta/agents/letta_agent_v2.py +++ b/letta/agents/letta_agent_v2.py @@ -1106,7 +1106,7 @@ class LettaAgentV2(BaseAgentV2): start_time = get_utc_timestamp_ns() agent_step_span.add_event(name="tool_execution_started") - sandbox_env_vars = {var.key: var.value for var in agent_state.tool_exec_environment_variables} + sandbox_env_vars = {var.key: var.value for var in agent_state.secrets} tool_execution_manager = ToolExecutionManager( agent_state=agent_state, message_manager=self.message_manager, diff --git a/letta/agents/voice_agent.py b/letta/agents/voice_agent.py index 642b9d61..e3c20acf 100644 --- a/letta/agents/voice_agent.py +++ b/letta/agents/voice_agent.py @@ -441,7 +441,7 @@ class VoiceAgent(BaseAgent): ) # Use ToolExecutionManager for modern tool execution - sandbox_env_vars = {var.key: var.value for var in agent_state.tool_exec_environment_variables} + sandbox_env_vars = {var.key: var.value for var in agent_state.secrets} tool_execution_manager = ToolExecutionManager( agent_state=agent_state, message_manager=self.message_manager, diff --git a/letta/orm/agent.py b/letta/orm/agent.py index 1a14adec..49c45982 100644 --- a/letta/orm/agent.py +++ b/letta/orm/agent.py @@ -233,6 +233,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin "identity_ids": [], "multi_agent_group": None, "tool_exec_environment_variables": [], + "secrets": [], } # Optional fields: only included if requested @@ -253,6 +254,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin "identity_ids": lambda: [i.id for i in self.identities], "multi_agent_group": lambda: self.multi_agent_group, "tool_exec_environment_variables": lambda: self.tool_exec_environment_variables, + "secrets": lambda: self.tool_exec_environment_variables, } include_relationships = set(optional_fields.keys() if include_relationships is None else include_relationships) @@ -324,6 +326,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin "identity_ids": [], "multi_agent_group": None, "tool_exec_environment_variables": [], + "secrets": [], } # Initialize include_relationships to an empty set if it's None @@ -344,7 +347,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin multi_agent_group = self.awaitable_attrs.multi_agent_group if "multi_agent_group" in include_relationships else none_async() tool_exec_environment_variables = ( self.awaitable_attrs.tool_exec_environment_variables - if "tool_exec_environment_variables" in include_relationships + if "tool_exec_environment_variables" in include_relationships or "secrets" in include_relationships else empty_list_async() ) file_agents = self.awaitable_attrs.file_agents if "memory" in include_relationships else empty_list_async() @@ -368,5 +371,6 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin state["identity_ids"] = [i.id for i in identities] state["multi_agent_group"] = multi_agent_group state["tool_exec_environment_variables"] = tool_exec_environment_variables + state["secrets"] = tool_exec_environment_variables return self.__pydantic_model__(**state) diff --git a/letta/schemas/agent.py b/letta/schemas/agent.py index cd00f54b..5ca23723 100644 --- a/letta/schemas/agent.py +++ b/letta/schemas/agent.py @@ -86,6 +86,11 @@ class AgentState(OrmMetadataBase, validate_assignment=True): sources: List[Source] = Field(..., description="The sources used by the agent.") tags: List[str] = Field(..., description="The tags associated with the agent.") tool_exec_environment_variables: List[AgentEnvironmentVariable] = Field( + default_factory=list, + description="Deprecated: use `secrets` field instead.", + deprecated=True, + ) + secrets: List[AgentEnvironmentVariable] = Field( default_factory=list, description="The environment variables for tool execution specific to this agent." ) project_id: Optional[str] = Field(None, description="The id of the project the agent belongs to.") @@ -133,7 +138,7 @@ class AgentState(OrmMetadataBase, validate_assignment=True): 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: + for agent_env_var_obj in self.secrets: per_agent_env_vars[agent_env_var_obj.key] = agent_env_var_obj.value return per_agent_env_vars @@ -222,9 +227,8 @@ class CreateAgent(BaseModel, validate_assignment=True): # deprecated=True, description="Deprecated: Project should now be passed via the X-Project header instead of in the request body. If using the sdk, this can be done via the new x_project field below.", ) - tool_exec_environment_variables: Optional[Dict[str, str]] = Field( - None, description="The environment variables for tool execution specific to this agent." - ) + tool_exec_environment_variables: Optional[Dict[str, str]] = Field(None, description="Deprecated: use `secrets` field instead.") + secrets: Optional[Dict[str, str]] = Field(None, description="The environment variables for tool execution specific to this agent.") memory_variables: Optional[Dict[str, str]] = Field(None, description="The variables that should be set for the agent.") project_id: Optional[str] = Field(None, description="The id of the project the agent belongs to.") template_id: Optional[str] = Field(None, description="The id of the template the agent belongs to.") @@ -328,9 +332,8 @@ class UpdateAgent(BaseModel): message_ids: Optional[List[str]] = Field(None, description="The ids of the messages in the agent's in-context memory.") description: Optional[str] = Field(None, description="The description of the agent.") metadata: Optional[Dict] = Field(None, description="The metadata of the agent.") - tool_exec_environment_variables: Optional[Dict[str, str]] = Field( - None, description="The environment variables for tool execution specific to this agent." - ) + tool_exec_environment_variables: Optional[Dict[str, str]] = Field(None, description="Deprecated: use `secrets` field instead") + secrets: Optional[Dict[str, str]] = Field(None, description="The environment variables for tool execution specific to this agent.") project_id: Optional[str] = Field(None, description="The id of the project the agent belongs to.") template_id: Optional[str] = Field(None, description="The id of the template the agent belongs to.") base_template_id: Optional[str] = Field(None, description="The base template id of the agent.") diff --git a/letta/serialize_schemas/marshmallow_agent.py b/letta/serialize_schemas/marshmallow_agent.py index fe861ba3..a8ecf589 100644 --- a/letta/serialize_schemas/marshmallow_agent.py +++ b/letta/serialize_schemas/marshmallow_agent.py @@ -40,6 +40,7 @@ class MarshmallowAgentSchema(BaseSchema): core_memory = fields.List(fields.Nested(SerializedBlockSchema)) tools = fields.List(fields.Nested(SerializedToolSchema)) tool_exec_environment_variables = fields.List(fields.Nested(SerializedAgentEnvironmentVariableSchema)) + secrets = fields.List(fields.Nested(SerializedAgentEnvironmentVariableSchema)) tags = fields.List(fields.Nested(SerializedAgentTagSchema)) def __init__(self, *args, session: sessionmaker, actor: User, max_steps: Optional[int] = None, **kwargs): @@ -214,6 +215,9 @@ class MarshmallowAgentSchema(BaseSchema): for env_var in data.get("tool_exec_environment_variables", []): # need to be re-set at load time env_var["value"] = "" + for env_var in data.get("secrets", []): + # need to be re-set at load time + env_var["value"] = "" return data @pre_load diff --git a/letta/services/agent_manager.py b/letta/services/agent_manager.py index ec35270e..e063477c 100644 --- a/letta/services/agent_manager.py +++ b/letta/services/agent_manager.py @@ -455,7 +455,8 @@ class AgentManager: [{"agent_id": aid, "identity_id": iid} for iid in identity_ids], ) - if agent_create.tool_exec_environment_variables: + agent_secrets = agent_create.secrets or agent_create.tool_exec_environment_variables + if agent_secrets: env_rows = [ { "agent_id": aid, @@ -463,7 +464,7 @@ class AgentManager: "value": val, "organization_id": actor.organization_id, } - for key, val in agent_create.tool_exec_environment_variables.items() + for key, val in agent_secrets.items() ] session.execute(insert(AgentEnvironmentVariable).values(env_rows)) @@ -674,7 +675,8 @@ class AgentManager: ) env_rows = [] - if agent_create.tool_exec_environment_variables: + agent_secrets = agent_create.secrets or agent_create.tool_exec_environment_variables + if agent_secrets: env_rows = [ { "agent_id": aid, @@ -682,7 +684,7 @@ class AgentManager: "value": val, "organization_id": actor.organization_id, } - for key, val in agent_create.tool_exec_environment_variables.items() + for key, val in agent_secrets.items() ] result = await session.execute(insert(AgentEnvironmentVariable).values(env_rows).returning(AgentEnvironmentVariable.id)) env_rows = [{**row, "id": env_var_id} for row, env_var_id in zip(env_rows, result.scalars().all())] @@ -701,8 +703,9 @@ class AgentManager: result = await new_agent.to_pydantic_async(include_relationships=include_relationships) - if agent_create.tool_exec_environment_variables and env_rows: + if agent_secrets and env_rows: result.tool_exec_environment_variables = [AgentEnvironmentVariable(**row) for row in env_rows] + result.secrets = [AgentEnvironmentVariable(**row) for row in env_rows] # initial message sequence (skip if _init_with_no_messages is True) if not _init_with_no_messages: @@ -894,7 +897,8 @@ class AgentManager: ) session.expire(agent, ["tags"]) - if agent_update.tool_exec_environment_variables is not None: + agent_secrets = agent_update.secrets or agent_update.tool_exec_environment_variables + if agent_secrets is not None: session.execute(delete(AgentEnvironmentVariable).where(AgentEnvironmentVariable.agent_id == aid)) env_rows = [ { @@ -903,7 +907,7 @@ class AgentManager: "value": v, "organization_id": agent.organization_id, } - for k, v in agent_update.tool_exec_environment_variables.items() + for k, v in agent_secrets.items() ] if env_rows: self._bulk_insert_pivot(session, AgentEnvironmentVariable.__table__, env_rows) @@ -1019,7 +1023,8 @@ class AgentManager: ) session.expire(agent, ["tags"]) - if agent_update.tool_exec_environment_variables is not None: + agent_secrets = agent_update.secrets or agent_update.tool_exec_environment_variables + if agent_secrets is not None: await session.execute(delete(AgentEnvironmentVariable).where(AgentEnvironmentVariable.agent_id == aid)) env_rows = [ { @@ -1028,7 +1033,7 @@ class AgentManager: "value": v, "organization_id": agent.organization_id, } - for k, v in agent_update.tool_exec_environment_variables.items() + for k, v in agent_secrets.items() ] if env_rows: await self._bulk_insert_pivot_async(session, AgentEnvironmentVariable.__table__, env_rows) @@ -1544,6 +1549,8 @@ class AgentManager: if env_vars: for var in agent.tool_exec_environment_variables: var.value = env_vars.get(var.key, "") + for var in agent.secrets: + var.value = env_vars.get(var.key, "") agent = agent.create(session, actor=actor) @@ -1627,6 +1634,7 @@ class AgentManager: # Remove stale variables stale_keys = set(existing_vars) - set(env_vars) agent.tool_exec_environment_variables = [var for var in updated_vars if var.key not in stale_keys] + agent.secrets = [var for var in updated_vars if var.key not in stale_keys] # Update the agent in the database agent.update(session, actor=actor) diff --git a/letta/services/agent_serialization_manager.py b/letta/services/agent_serialization_manager.py index dbfb3250..01ded468 100644 --- a/letta/services/agent_serialization_manager.py +++ b/letta/services/agent_serialization_manager.py @@ -209,8 +209,10 @@ class AgentSerializationManager: agent_schema.id = agent_file_id # wipe the values of tool_exec_environment_variables (they contain secrets) - if agent_schema.tool_exec_environment_variables: - agent_schema.tool_exec_environment_variables = {key: "" for key in agent_schema.tool_exec_environment_variables} + agent_secrets = agent_schema.secrets or agent_schema.tool_exec_environment_variables + if agent_secrets: + agent_schema.tool_exec_environment_variables = {key: "" for key in agent_secrets} + agent_schema.secrets = {key: "" for key in agent_secrets} if agent_schema.messages: for message in agent_schema.messages: @@ -655,10 +657,16 @@ class AgentSerializationManager: if agent_data.get("source_ids"): agent_data["source_ids"] = [file_to_db_ids[file_id] for file_id in agent_data["source_ids"]] - if env_vars and agent_data.get("tool_exec_environment_variables"): + if env_vars and agent_data.get("secrets"): # update environment variable values from the provided env_vars dict + for key in agent_data["secrets"]: + agent_data["secrets"][key] = env_vars.get(key, "") + agent_data["tool_exec_environment_variables"][key] = env_vars.get(key, "") + elif env_vars and agent_data.get("tool_exec_environment_variables"): + # also handle tool_exec_environment_variables for backwards compatibility for key in agent_data["tool_exec_environment_variables"]: agent_data["tool_exec_environment_variables"][key] = env_vars.get(key, "") + agent_data["secrets"][key] = env_vars.get(key, "") # Override project_id if provided if project_id: diff --git a/letta/services/tool_executor/composio_tool_executor.py b/letta/services/tool_executor/composio_tool_executor.py index d2e8e64e..30030f3b 100644 --- a/letta/services/tool_executor/composio_tool_executor.py +++ b/letta/services/tool_executor/composio_tool_executor.py @@ -51,7 +51,7 @@ class ExternalComposioToolExecutor(ToolExecutor): def _get_entity_id(self, agent_state: AgentState) -> Optional[str]: """Extract the entity ID from environment variables.""" - for env_var in agent_state.tool_exec_environment_variables: + for env_var in agent_state.secrets: if env_var.key == COMPOSIO_ENTITY_ENV_VAR_KEY: return env_var.value return None diff --git a/tests/helpers/utils.py b/tests/helpers/utils.py index 72df6806..cd43061e 100644 --- a/tests/helpers/utils.py +++ b/tests/helpers/utils.py @@ -108,11 +108,16 @@ def comprehensive_agent_checks(agent: AgentState, request: Union[CreateAgent, Up assert agent.metadata == request.metadata, f"Metadata mismatch: {agent.metadata} != {request.metadata}" # Assert agent env vars - if hasattr(request, "tool_exec_environment_variables"): + if hasattr(request, "tool_exec_environment_variables") and request.tool_exec_environment_variables: for agent_env_var in agent.tool_exec_environment_variables: assert agent_env_var.key in request.tool_exec_environment_variables assert request.tool_exec_environment_variables[agent_env_var.key] == agent_env_var.value assert agent_env_var.organization_id == actor.organization_id + if hasattr(request, "secrets") and request.secrets: + for agent_env_var in agent.secrets: + assert agent_env_var.key in request.secrets + assert request.secrets[agent_env_var.key] == agent_env_var.value + assert agent_env_var.organization_id == actor.organization_id # Assert agent type if hasattr(request, "agent_type"): diff --git a/tests/integration_test_async_tool_sandbox.py b/tests/integration_test_async_tool_sandbox.py index a746f4c1..cf729e5d 100644 --- a/tests/integration_test_async_tool_sandbox.py +++ b/tests/integration_test_async_tool_sandbox.py @@ -589,7 +589,7 @@ async def test_local_sandbox_per_agent_env(disable_e2b_api_key, get_env_tool, ag manager.create_sandbox_env_var(SandboxEnvironmentVariableCreate(key=key, value=wrong_val), sandbox_config_id=config.id, actor=test_user) correct_val = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(20)) - agent_state.tool_exec_environment_variables = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)] + agent_state.secrets = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)] sandbox = AsyncToolSandboxLocal(get_env_tool.name, {}, user=test_user) result = await sandbox.run(agent_state=agent_state) @@ -812,7 +812,7 @@ async def test_e2b_sandbox_per_agent_env(check_e2b_key_is_set, get_env_tool, age actor=test_user, ) - agent_state.tool_exec_environment_variables = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)] + agent_state.secrets = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)] sandbox = AsyncToolSandboxE2B(get_env_tool.name, {}, user=test_user) result = await sandbox.run(agent_state=agent_state) @@ -1116,7 +1116,7 @@ async def test_local_sandbox_async_per_agent_env(disable_e2b_api_key, async_get_ manager.create_sandbox_env_var(SandboxEnvironmentVariableCreate(key=key, value=wrong_val), sandbox_config_id=config.id, actor=test_user) correct_val = "correct_async_local_value" - agent_state.tool_exec_environment_variables = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)] + agent_state.secrets = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)] sandbox = AsyncToolSandboxLocal(async_get_env_tool.name, {}, user=test_user) result = await sandbox.run(agent_state=agent_state) @@ -1141,7 +1141,7 @@ async def test_e2b_sandbox_async_per_agent_env(check_e2b_key_is_set, async_get_e actor=test_user, ) - agent_state.tool_exec_environment_variables = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)] + agent_state.secrets = [AgentEnvironmentVariable(key=key, value=correct_val, agent_id=agent_state.id)] sandbox = AsyncToolSandboxE2B(async_get_env_tool.name, {}, user=test_user) result = await sandbox.run(agent_state=agent_state) diff --git a/tests/integration_test_tool_execution_sandbox.py b/tests/integration_test_tool_execution_sandbox.py index 2cdee949..cd3dc394 100644 --- a/tests/integration_test_tool_execution_sandbox.py +++ b/tests/integration_test_tool_execution_sandbox.py @@ -359,9 +359,7 @@ def test_local_sandbox_per_agent_env(disable_e2b_api_key, get_env_tool, agent_st # 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) - ] + agent_state.secrets = [AgentEnvironmentVariable(key=key, value=correct_long_random_string, agent_id=agent_state.id)] # Create tool and args args = {} @@ -569,9 +567,7 @@ def test_e2b_sandbox_per_agent_env(check_e2b_key_is_set, get_env_tool, agent_sta # 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) - ] + agent_state.secrets = [AgentEnvironmentVariable(key=key, value=correct_long_random_string, agent_id=agent_state.id)] # Create tool and args args = {} diff --git a/tests/test_agent_serialization.py b/tests/test_agent_serialization.py index 90588f0c..e5d5b78c 100644 --- a/tests/test_agent_serialization.py +++ b/tests/test_agent_serialization.py @@ -298,7 +298,7 @@ def _compare_agent_state_model_dump(d1: Dict[str, Any], d2: Dict[str, Any], log: if not isinstance(v1, list) or not isinstance(v2, list) or len(v1) != len(v2): _log_mismatch(key, v1, v2, log) return False - elif key == "tool_exec_environment_variables": + elif key == "tool_exec_environment_variables" or key == "secrets": if not isinstance(v1, list) or not isinstance(v2, list) or len(v1) != len(v2): _log_mismatch(key, v1, v2, log) return False diff --git a/tests/test_agent_serialization_v2.py b/tests/test_agent_serialization_v2.py index cfff36da..768e4545 100644 --- a/tests/test_agent_serialization_v2.py +++ b/tests/test_agent_serialization_v2.py @@ -503,7 +503,7 @@ def _compare_agents(orig: AgentSchema, imp: AgentSchema, index: int) -> List[str errors.append(f"Agent {index}: tool_rules mismatch") # Environment variables - if orig.tool_exec_environment_variables != imp.tool_exec_environment_variables: + if orig.secrets != imp.secrets: errors.append(f"Agent {index}: tool_exec_environment_variables mismatch") # Messages @@ -1084,16 +1084,16 @@ class TestAgentFileExport: exported_agent = agent_file.agents[0] # verify environment variables exist but values are scrubbed (empty strings) - assert exported_agent.tool_exec_environment_variables is not None - assert len(exported_agent.tool_exec_environment_variables) == 3 - assert "API_KEY" in exported_agent.tool_exec_environment_variables - assert "DATABASE_PASSWORD" in exported_agent.tool_exec_environment_variables - assert "TOKEN" in exported_agent.tool_exec_environment_variables + assert exported_agent.secrets is not None + assert len(exported_agent.secrets) == 3 + assert "API_KEY" in exported_agent.secrets + assert "DATABASE_PASSWORD" in exported_agent.secrets + assert "TOKEN" in exported_agent.secrets # most importantly: verify all secret values have been wiped - assert exported_agent.tool_exec_environment_variables["API_KEY"] == "" - assert exported_agent.tool_exec_environment_variables["DATABASE_PASSWORD"] == "" - assert exported_agent.tool_exec_environment_variables["TOKEN"] == "" + assert exported_agent.secrets["API_KEY"] == "" + assert exported_agent.secrets["DATABASE_PASSWORD"] == "" + assert exported_agent.secrets["TOKEN"] == "" # verify no secret values appear anywhere in the exported data assert "super-secret-api-key-12345" not in str(agent_file) @@ -1287,9 +1287,9 @@ class TestAgentFileImport: # verify values are scrubbed in export exported_agent = agent_file.agents[0] - assert exported_agent.tool_exec_environment_variables["API_KEY"] == "" - assert exported_agent.tool_exec_environment_variables["DATABASE_URL"] == "" - assert exported_agent.tool_exec_environment_variables["SECRET_TOKEN"] == "" + assert exported_agent.secrets["API_KEY"] == "" + assert exported_agent.secrets["DATABASE_URL"] == "" + assert exported_agent.secrets["SECRET_TOKEN"] == "" # import with new environment variable values new_env_vars = {