letta-nightly 0.12.1.dev20251024104217__py3-none-any.whl → 0.13.0.dev20251024223017__py3-none-any.whl

letta/services/sandbox_config_manager.py

@@ -5,7 +5,7 @@ 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 SandboxType
+from letta.schemas.enums import PrimitiveType, SandboxType
 from letta.schemas.environment_variables import (
     SandboxEnvironmentVariable as PydanticEnvVar,
     SandboxEnvironmentVariableCreate,
@@ -20,6 +20,7 @@ from letta.schemas.sandbox_config import (
 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__)
 
@@ -101,6 +102,7 @@ class SandboxConfigManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
     async def update_sandbox_config_async(
         self, sandbox_config_id: str, sandbox_update: SandboxConfigUpdate, actor: PydanticUser
     ) -> PydanticSandboxConfig:
@@ -129,6 +131,7 @@ class SandboxConfigManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
     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:
@@ -176,6 +179,7 @@ class SandboxConfigManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
     async def create_sandbox_env_var_async(
         self, env_var_create: SandboxEnvironmentVariableCreate, sandbox_config_id: str, actor: PydanticUser
     ) -> PydanticEnvVar:
@@ -198,6 +202,12 @@ class SandboxConfigManager:
             return db_env_var
         else:
             async with db_registry.async_session() as session:
+                # Explicitly encrypt the value before storing
+                from letta.schemas.secret import Secret
+
+                if env_var.value is not None:
+                    env_var.value_enc = Secret.from_plaintext(env_var.value)
+
                 env_var = SandboxEnvVarModel(**env_var.model_dump(to_orm=True, exclude_none=True))
                 await env_var.create_async(session, actor=actor)
                 return env_var.to_pydantic()
@@ -211,6 +221,31 @@ class SandboxConfigManager:
         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
+                existing_value = None
+                if env_var.value_enc:
+                    existing_secret = Secret.from_encrypted(env_var.value_enc)
+                    existing_value = existing_secret.get_plaintext()
+                elif env_var.value:
+                    existing_value = env_var.value
+
+                # Only re-encrypt if different
+                if existing_value != update_data["value"]:
+                    env_var.value_enc = Secret.from_plaintext(update_data["value"]).get_encrypted()
+                    # Keep plaintext for dual-write during migration
+                    env_var.value = update_data["value"]
+
+                # 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:
@@ -235,6 +270,7 @@ class SandboxConfigManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
     async def list_sandbox_env_vars_async(
         self,
         sandbox_config_id: str,
@@ -271,25 +307,31 @@ class SandboxConfigManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
     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:
-            result[env_var.key] = env_var.value
+            # Decrypt the value before returning
+            value_secret = env_var.get_value_secret()
+            result[env_var.key] = value_secret.get_plaintext()
         return result
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
     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)
-        return {env_var.key: env_var.value for env_var in env_vars}
+        # Decrypt values before returning
+        return {env_var.key: env_var.get_value_secret().get_plaintext() for env_var in env_vars}
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="sandbox_config_id", expected_prefix=PrimitiveType.SANDBOX_CONFIG)
     async def get_sandbox_env_var_by_key_and_sandbox_config_id_async(
         self, key: str, sandbox_config_id: str, actor: Optional[PydanticUser] = None
     ) -> Optional[PydanticEnvVar]:

letta/services/source_manager.py

@@ -11,11 +11,12 @@ from letta.orm.source import Source as SourceModel
 from letta.orm.sources_agents import SourcesAgents
 from letta.otel.tracing import trace_method
 from letta.schemas.agent import AgentState as PydanticAgentState
-from letta.schemas.enums import VectorDBProvider
+from letta.schemas.enums import PrimitiveType, VectorDBProvider
 from letta.schemas.source import Source as PydanticSource, SourceUpdate
 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
 
 
 class SourceManager:
@@ -61,7 +62,11 @@ class SourceManager:
     @trace_method
     async def create_source(self, source: PydanticSource, actor: PydanticUser) -> PydanticSource:
         """Create a new source based on the PydanticSource schema."""
-        db_source = await self.get_source_by_id(source.id, actor=actor)
+        try:
+            db_source = await self.get_source_by_id(source.id, actor=actor)
+        except NoResultFound:
+            db_source = None
+
         if db_source:
             return db_source
         else:
@@ -197,6 +202,7 @@ class SourceManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE)
     async def update_source(self, source_id: str, source_update: SourceUpdate, actor: PydanticUser) -> PydanticSource:
         """Update a source by its ID with the given SourceUpdate object."""
         async with db_registry.async_session() as session:
@@ -220,6 +226,7 @@ class SourceManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE)
     async def delete_source(self, source_id: str, actor: PydanticUser) -> PydanticSource:
         """Delete a source by its ID."""
         async with db_registry.async_session() as session:
@@ -264,6 +271,7 @@ class SourceManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE)
     async def list_attached_agents(
         self, source_id: str, actor: PydanticUser, ids_only: bool = False
     ) -> Union[List[PydanticAgentState], List[str]]:
@@ -317,6 +325,7 @@ class SourceManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE)
     async def get_agents_for_source_id(self, source_id: str, actor: PydanticUser) -> List[str]:
         """
         Get all agent IDs associated with a given source ID.
@@ -343,14 +352,12 @@ class SourceManager:
     # TODO: We make actor optional for now, but should most likely be enforced due to security reasons
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE)
     async def get_source_by_id(self, source_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticSource]:
         """Retrieve a source by its ID."""
         async with db_registry.async_session() as session:
-            try:
-                source = await SourceModel.read_async(db_session=session, identifier=source_id, actor=actor)
-                return source.to_pydantic()
-            except NoResultFound:
-                return None
+            source = await SourceModel.read_async(db_session=session, identifier=source_id, actor=actor)
+            return source.to_pydantic()
 
     @enforce_types
     @trace_method
@@ -364,7 +371,7 @@ class SourceManager:
                 limit=1,
             )
             if not sources:
-                return None
+                raise NoResultFound(f"Source with name={source_name} not found.")
             else:
                 return sources[0].to_pydantic()
 

letta/services/step_manager.py

@@ -13,7 +13,7 @@ from letta.orm.sqlalchemy_base import AccessType
 from letta.orm.step import Step as StepModel
 from letta.orm.step_metrics import StepMetrics as StepMetricsModel
 from letta.otel.tracing import get_trace_id, trace_method
-from letta.schemas.enums import StepStatus
+from letta.schemas.enums import PrimitiveType, StepStatus
 from letta.schemas.letta_stop_reason import LettaStopReason, StopReasonType
 from letta.schemas.message import Message as PydanticMessage
 from letta.schemas.openai.chat_completion_response import UsageStatistics
@@ -22,6 +22,7 @@ from letta.schemas.step_metrics import StepMetrics as PydanticStepMetrics
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
 from letta.utils import enforce_types
+from letta.validators import raise_on_invalid_id
 
 
 class FeedbackType(str, Enum):
@@ -32,6 +33,8 @@ class FeedbackType(str, Enum):
 class StepManager:
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT)
+    @raise_on_invalid_id(param_name="run_id", expected_prefix=PrimitiveType.RUN)
     async def list_steps_async(
         self,
         actor: PydanticUser,
@@ -79,6 +82,10 @@ class StepManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT)
+    @raise_on_invalid_id(param_name="provider_id", expected_prefix=PrimitiveType.PROVIDER)
+    @raise_on_invalid_id(param_name="run_id", expected_prefix=PrimitiveType.RUN)
+    @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP)
     def log_step(
         self,
         actor: PydanticUser,
@@ -133,6 +140,10 @@ class StepManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT)
+    @raise_on_invalid_id(param_name="provider_id", expected_prefix=PrimitiveType.PROVIDER)
+    @raise_on_invalid_id(param_name="run_id", expected_prefix=PrimitiveType.RUN)
+    @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP)
     async def log_step_async(
         self,
         actor: PydanticUser,
@@ -196,6 +207,7 @@ class StepManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP)
     async def get_step_async(self, step_id: str, actor: PydanticUser) -> PydanticStep:
         async with db_registry.async_session() as session:
             step = await StepModel.read_async(db_session=session, identifier=step_id, actor=actor)
@@ -203,6 +215,7 @@ class StepManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP)
     async def get_step_metrics_async(self, step_id: str, actor: PydanticUser) -> PydanticStepMetrics:
         async with db_registry.async_session() as session:
             metrics = await StepMetricsModel.read_async(db_session=session, identifier=step_id, actor=actor)
@@ -210,6 +223,7 @@ class StepManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP)
     async def add_feedback_async(
         self, step_id: str, feedback: FeedbackType | None, actor: PydanticUser, tags: list[str] | None = None
     ) -> PydanticStep:
@@ -225,6 +239,7 @@ class StepManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP)
     async def update_step_transaction_id(self, actor: PydanticUser, step_id: str, transaction_id: str) -> PydanticStep:
         """Update the transaction ID for a step.
 
@@ -252,6 +267,7 @@ class StepManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP)
     async def list_step_messages_async(
         self,
         step_id: str,
@@ -276,6 +292,7 @@ class StepManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP)
     async def update_step_stop_reason(self, actor: PydanticUser, step_id: str, stop_reason: StopReasonType) -> PydanticStep:
         """Update the stop reason for a step.
 
@@ -303,6 +320,7 @@ class StepManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP)
     async def update_step_error_async(
         self,
         actor: PydanticUser,
@@ -348,6 +366,7 @@ class StepManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP)
     async def update_step_success_async(
         self,
         actor: PydanticUser,
@@ -388,6 +407,7 @@ class StepManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP)
     async def update_step_cancelled_async(
         self,
         actor: PydanticUser,
@@ -423,6 +443,9 @@ class StepManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP)
+    @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT)
+    @raise_on_invalid_id(param_name="run_id", expected_prefix=PrimitiveType.RUN)
     async def record_step_metrics_async(
         self,
         actor: PydanticUser,