diff --git a/letta/agent.py b/letta/agent.py index 0dde8fef..bf9e5e7a 100644 --- a/letta/agent.py +++ b/letta/agent.py @@ -1,6 +1,5 @@ import asyncio import json -import os import time import traceback import warnings @@ -71,7 +70,7 @@ from letta.services.step_manager import StepManager from letta.services.telemetry_manager import NoopTelemetryManager, TelemetryManager from letta.services.tool_executor.tool_execution_sandbox import ToolExecutionSandbox from letta.services.tool_manager import ToolManager -from letta.settings import settings, summarizer_settings +from letta.settings import model_settings, settings, summarizer_settings from letta.streaming_interface import StreamingRefreshCLIInterface from letta.system import get_heartbeat, get_token_limit_warning, package_function_response, package_summarize_message, package_user_message from letta.utils import count_tokens, get_friendly_error_msg, get_tool_call_id, log_telemetry, parse_json, validate_function_response @@ -1304,7 +1303,7 @@ class Agent(BaseAgent): ) async def get_context_window_async(self) -> ContextWindowOverview: - if os.getenv("LETTA_ENVIRONMENT") == "PRODUCTION" and os.getenv("ANTHROPIC_API_KEY"): + if settings.environment == "PRODUCTION" and model_settings.anthropic_api_key: return await self.get_context_window_from_anthropic_async() return await self.get_context_window_from_tiktoken_async() diff --git a/letta/functions/function_sets/multi_agent.py b/letta/functions/function_sets/multi_agent.py index 373512d2..17bd5586 100644 --- a/letta/functions/function_sets/multi_agent.py +++ b/letta/functions/function_sets/multi_agent.py @@ -1,6 +1,5 @@ import asyncio import json -import os from concurrent.futures import ThreadPoolExecutor, as_completed from typing import TYPE_CHECKING, List @@ -138,7 +137,7 @@ def send_message_to_agent_async(self: "Agent", message: str, other_agent_id: str Returns: str: A confirmation message indicating the message was successfully sent. """ - if os.getenv("LETTA_ENVIRONMENT") == "PRODUCTION": + if settings.environment == "PRODUCTION": raise RuntimeError("This tool is not allowed to be run on Letta Cloud.") message = ( diff --git a/letta/otel/resource.py b/letta/otel/resource.py index d11e3104..f8724e75 100644 --- a/letta/otel/resource.py +++ b/letta/otel/resource.py @@ -1,16 +1,16 @@ -import os import sys import uuid from opentelemetry.sdk.resources import Resource from letta import __version__ as letta_version +from letta.settings import settings _resources = {} def get_resource(service_name: str) -> Resource: - _env = os.getenv("LETTA_ENVIRONMENT") + _env = settings.environment if service_name not in _resources: resource_dict = { "service.name": service_name, diff --git a/letta/services/agent_manager.py b/letta/services/agent_manager.py index cd4f180f..46d25e16 100644 --- a/letta/services/agent_manager.py +++ b/letta/services/agent_manager.py @@ -1,5 +1,4 @@ import asyncio -import os from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Set, Tuple @@ -3423,7 +3422,7 @@ class AgentManager: ) calculator = ContextWindowCalculator() - if os.getenv("LETTA_ENVIRONMENT") == "PRODUCTION" or agent_state.llm_config.model_endpoint_type == "anthropic": + if settings.environment == "PRODUCTION" or agent_state.llm_config.model_endpoint_type == "anthropic": anthropic_client = LLMClient.create(provider_type=ProviderType.anthropic, actor=actor) model = agent_state.llm_config.model if agent_state.llm_config.model_endpoint_type == "anthropic" else None diff --git a/letta/services/helpers/agent_manager_helper.py b/letta/services/helpers/agent_manager_helper.py index 5e147c06..b302039a 100644 --- a/letta/services/helpers/agent_manager_helper.py +++ b/letta/services/helpers/agent_manager_helper.py @@ -1,4 +1,3 @@ -import os import uuid from datetime import datetime from typing import List, Literal, Optional, Set @@ -1208,7 +1207,7 @@ def calculate_base_tools(is_v2: bool) -> Set[str]: def calculate_multi_agent_tools() -> Set[str]: """Calculate multi-agent tools, excluding local-only tools in production environment.""" - if os.getenv("LETTA_ENVIRONMENT") == "PRODUCTION": + if settings.environment == "PRODUCTION": return set(MULTI_AGENT_TOOLS) - set(LOCAL_ONLY_MULTI_AGENT_TOOLS) else: return set(MULTI_AGENT_TOOLS) diff --git a/letta/services/tool_executor/multi_agent_tool_executor.py b/letta/services/tool_executor/multi_agent_tool_executor.py index 6e7e2d7b..7aa57bae 100644 --- a/letta/services/tool_executor/multi_agent_tool_executor.py +++ b/letta/services/tool_executor/multi_agent_tool_executor.py @@ -1,5 +1,4 @@ import asyncio -import os from typing import Any, Dict, List, Optional from letta.log import get_logger @@ -13,6 +12,7 @@ from letta.schemas.tool import Tool from letta.schemas.tool_execution_result import ToolExecutionResult from letta.schemas.user import User from letta.services.tool_executor.tool_executor_base import ToolExecutor +from letta.settings import settings logger = get_logger(__name__) @@ -112,7 +112,7 @@ class LettaMultiAgentToolExecutor(ToolExecutor): } async def send_message_to_agent_async(self, agent_state: AgentState, message: str, other_agent_id: str) -> str: - if os.getenv("LETTA_ENVIRONMENT") == "PRODUCTION": + if settings.environment == "PRODUCTION": raise RuntimeError("This tool is not allowed to be run on Letta Cloud.") # 1) Build the prefixed system‐message diff --git a/letta/services/tool_manager.py b/letta/services/tool_manager.py index d990ff92..d85e5c5c 100644 --- a/letta/services/tool_manager.py +++ b/letta/services/tool_manager.py @@ -1,5 +1,4 @@ import importlib -import os import warnings from typing import List, Optional, Set, Union @@ -331,9 +330,7 @@ class ToolManager: # TODO: This requires a deeper rethink about how we keep all our internal tools up-to-date if not after and upsert_base_tools: existing_tool_names = {tool.name for tool in tools} - base_tool_names = ( - LETTA_TOOL_SET - set(LOCAL_ONLY_MULTI_AGENT_TOOLS) if os.getenv("LETTA_ENVIRONMENT") == "PRODUCTION" else LETTA_TOOL_SET - ) + base_tool_names = LETTA_TOOL_SET - set(LOCAL_ONLY_MULTI_AGENT_TOOLS) if settings.environment == "PRODUCTION" else LETTA_TOOL_SET missing_base_tools = base_tool_names - existing_tool_names # If any base tools are missing, upsert all base tools diff --git a/letta/settings.py b/letta/settings.py index 2e052a3d..87a87534 100644 --- a/letta/settings.py +++ b/letta/settings.py @@ -205,6 +205,7 @@ class Settings(BaseSettings): letta_dir: Optional[Path] = Field(Path.home() / ".letta", alias="LETTA_DIR") debug: Optional[bool] = False cors_origins: Optional[list] = cors_origins + environment: Optional[str] = Field(default=None, description="Application environment (PRODUCTION, DEV, etc.)") # SSE Streaming keepalive settings enable_keepalive: bool = Field(True, description="Enable keepalive messages in SSE streams to prevent timeouts") diff --git a/tests/test_managers.py b/tests/test_managers.py index 2ccfdd49..fd2f32f0 100644 --- a/tests/test_managers.py +++ b/tests/test_managers.py @@ -101,7 +101,7 @@ from letta.server.server import SyncServer from letta.services.block_manager import BlockManager from letta.services.helpers.agent_manager_helper import calculate_base_tools, calculate_multi_agent_tools, validate_agent_exists_async from letta.services.step_manager import FeedbackType -from letta.settings import tool_settings +from letta.settings import settings, tool_settings from letta.utils import calculate_file_defaults_based_on_context_window from tests.helpers.utils import comprehensive_agent_checks, validate_context_window_overview from tests.utils import random_string @@ -862,7 +862,7 @@ def test_calculate_multi_agent_tools(set_letta_environment): """Test that calculate_multi_agent_tools excludes local-only tools in production.""" result = calculate_multi_agent_tools() - if set_letta_environment == "PRODUCTION": + if settings.environment == "PRODUCTION": # Production environment should exclude local-only tools expected_tools = set(MULTI_AGENT_TOOLS) - set(LOCAL_ONLY_MULTI_AGENT_TOOLS) assert result == expected_tools, "Production should exclude local-only multi-agent tools" @@ -889,7 +889,7 @@ async def test_upsert_base_tools_excludes_local_only_in_production(server: SyncS tools = await server.tool_manager.upsert_base_tools_async(actor=default_user) tool_names = {tool.name for tool in tools} - if set_letta_environment == "PRODUCTION": + if settings.environment == "PRODUCTION": # Production environment should exclude local-only multi-agent tools for local_only_tool in LOCAL_ONLY_MULTI_AGENT_TOOLS: assert local_only_tool not in tool_names, f"Local-only tool '{local_only_tool}' should not be upserted in production" @@ -912,7 +912,7 @@ async def test_upsert_multi_agent_tools_only(server: SyncServer, default_user, s tools = await server.tool_manager.upsert_base_tools_async(actor=default_user, allowed_types={ToolType.LETTA_MULTI_AGENT_CORE}) tool_names = {tool.name for tool in tools} - if set_letta_environment == "PRODUCTION": + if settings.environment == "PRODUCTION": # Should only have non-local multi-agent tools expected_tools = set(MULTI_AGENT_TOOLS) - set(LOCAL_ONLY_MULTI_AGENT_TOOLS) assert tool_names == expected_tools, "Production multi-agent upsert should exclude local-only tools" @@ -991,16 +991,14 @@ async def test_create_agent_with_default_source(server: SyncServer, default_user server.agent_manager.delete_agent(created_agent_no_source.id, default_user) -@pytest.fixture(params=["", "PRODUCTION"]) -def set_letta_environment(request): - original = os.environ.get("LETTA_ENVIRONMENT") - os.environ["LETTA_ENVIRONMENT"] = request.param +@pytest.fixture(params=[None, "PRODUCTION"]) +def set_letta_environment(request, monkeypatch): + # Patch the settings.environment attribute + original = settings.environment + monkeypatch.setattr(settings, "environment", request.param) yield request.param - # Restore original environment variable - if original is not None: - os.environ["LETTA_ENVIRONMENT"] = original - else: - os.environ.pop("LETTA_ENVIRONMENT", None) + # Restore original environment + monkeypatch.setattr(settings, "environment", original) async def test_get_context_window_basic( @@ -3801,7 +3799,7 @@ async def test_upsert_base_tools(server: SyncServer, default_user): tools = await server.tool_manager.upsert_base_tools_async(actor=default_user) # Calculate expected tools accounting for production filtering - if os.getenv("LETTA_ENVIRONMENT") == "PRODUCTION": + if settings.environment == "PRODUCTION": expected_tool_names = sorted(LETTA_TOOL_SET - set(LOCAL_ONLY_MULTI_AGENT_TOOLS)) else: expected_tool_names = sorted(LETTA_TOOL_SET) @@ -3853,7 +3851,7 @@ async def test_upsert_filtered_base_tools(server: SyncServer, default_user, tool tool_names = sorted([t.name for t in tools]) # Adjust expected names for multi-agent tools in production - if tool_type == ToolType.LETTA_MULTI_AGENT_CORE and os.getenv("LETTA_ENVIRONMENT") == "PRODUCTION": + if tool_type == ToolType.LETTA_MULTI_AGENT_CORE and settings.environment == "PRODUCTION": expected_sorted = sorted(set(expected_names) - set(LOCAL_ONLY_MULTI_AGENT_TOOLS)) else: expected_sorted = sorted(expected_names)