From 591420876addd875302dae54d691b79f2b5c9a25 Mon Sep 17 00:00:00 2001 From: jnjpng Date: Tue, 16 Dec 2025 16:10:24 -0800 Subject: [PATCH] fix: correct decorator order for trace_method and raise_on_invalid_id (#7226) Swap the order of @trace_method and @raise_on_invalid_id decorators across all service managers so that @trace_method is always the first wrapper applied to the function (positioned directly above the method). This ensures the ID validation happens before tracing begins, which is the intended execution order. Files modified: - agent_manager.py (23 occurrences) - archive_manager.py (11 occurrences) - block_manager.py (7 occurrences) - file_manager.py (6 occurrences) - group_manager.py (9 occurrences) - identity_manager.py (10 occurrences) - job_manager.py (7 occurrences) - message_manager.py (2 occurrences) - provider_manager.py (3 occurrences) - sandbox_config_manager.py (7 occurrences) - source_manager.py (5 occurrences) - step_manager.py (13 occurrences) --- letta/services/agent_manager.py | 44 ++++++++++++------------ letta/services/archive_manager.py | 22 ++++++------ letta/services/block_manager.py | 14 ++++---- letta/services/file_manager.py | 12 +++---- letta/services/group_manager.py | 18 +++++----- letta/services/identity_manager.py | 20 +++++------ letta/services/job_manager.py | 14 ++++---- letta/services/message_manager.py | 4 +-- letta/services/provider_manager.py | 6 ++-- letta/services/sandbox_config_manager.py | 14 ++++---- letta/services/source_manager.py | 10 +++--- letta/services/step_manager.py | 26 +++++++------- 12 files changed, 102 insertions(+), 102 deletions(-) diff --git a/letta/services/agent_manager.py b/letta/services/agent_manager.py index f7f6702c..d0dd21a0 100644 --- a/letta/services/agent_manager.py +++ b/letta/services/agent_manager.py @@ -717,8 +717,8 @@ class AgentManager: return await self.append_to_in_context_messages_async(init_messages, agent_id=agent_state.id, actor=actor) @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def update_agent_async( self, agent_id: str, @@ -1081,8 +1081,8 @@ class AgentManager: return await AgentModel.size_async(db_session=session, actor=actor) @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def get_agent_by_id_async( self, agent_id: str, @@ -1142,8 +1142,8 @@ class AgentManager: raise @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def get_agent_archive_ids_async(self, agent_id: str, actor: PydanticUser) -> List[str]: """Get all archive IDs associated with an agent.""" from letta.orm import ArchivesAgents @@ -1156,8 +1156,8 @@ class AgentManager: return archive_ids @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def validate_agent_exists_async(self, agent_id: str, actor: PydanticUser) -> None: """ Validate that an agent exists and user has access to it. @@ -1174,8 +1174,8 @@ class AgentManager: await validate_agent_exists_async(session, agent_id, actor) @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def delete_agent_async(self, agent_id: str, actor: PydanticUser) -> None: """ Deletes an agent and its associated relationships. @@ -1237,8 +1237,8 @@ class AgentManager: # TODO: This can be fixed by having an actual relationship in the ORM for message_ids # TODO: This can also be made more efficient, instead of getting, setting, we can do it all in one db session for one query. @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def get_in_context_messages(self, agent_id: str, actor: PydanticUser) -> List[PydanticMessage]: agent_state = await self.get_agent_by_id_async(agent_id=agent_id, actor=actor) return await self.message_manager.get_messages_by_ids_async(message_ids=agent_state.message_ids, actor=actor) @@ -1250,8 +1250,8 @@ class AgentManager: return self.message_manager.get_message_by_id(message_id=message_ids[0], actor=actor) @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def get_system_message_async(self, agent_id: str, actor: PydanticUser) -> PydanticMessage: agent = await self.get_agent_by_id_async(agent_id=agent_id, include_relationships=[], actor=actor) return await self.message_manager.get_message_by_id_async(message_id=agent.message_ids[0], actor=actor) @@ -1432,8 +1432,8 @@ class AgentManager: return self.update_agent(agent_id=agent_id, agent_update=UpdateAgent(message_ids=message_ids), actor=actor) @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def set_in_context_messages_async(self, agent_id: str, message_ids: List[str], actor: PydanticUser) -> PydanticAgentState: return await self.update_agent_async(agent_id=agent_id, agent_update=UpdateAgent(message_ids=message_ids), actor=actor) @@ -1543,8 +1543,8 @@ class AgentManager: return agent_state @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def update_memory_if_changed_async(self, agent_id: str, new_memory: Memory, actor: PydanticUser) -> PydanticAgentState: """ Update internal memory object and system prompt if there have been modifications. @@ -1656,9 +1656,9 @@ class AgentManager: # Source Management # ====================================================================================================================== @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE) + @trace_method async def attach_source_async(self, agent_id: str, source_id: str, actor: PydanticUser) -> PydanticAgentState: """ Attaches a source to an agent. @@ -1732,8 +1732,8 @@ class AgentManager: self.append_to_in_context_messages(messages=[message], agent_id=agent_id, actor=actor) @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def append_system_message_async(self, agent_id: str, content: str, actor: PydanticUser): """ Async version of append_system_message. @@ -1820,9 +1820,9 @@ class AgentManager: return [source.to_pydantic() for source in sources] @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE) + @trace_method async def detach_source_async(self, agent_id: str, source_id: str, actor: PydanticUser) -> PydanticAgentState: """ Detaches a source from an agent. @@ -1909,9 +1909,9 @@ class AgentManager: return block.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK) + @trace_method async def attach_block_async(self, agent_id: str, block_id: str, actor: PydanticUser) -> PydanticAgentState: """Attaches a block to an agent. For sleeptime agents, also attaches to paired agents in the same group.""" async with db_registry.async_session() as session: @@ -2503,9 +2503,9 @@ class AgentManager: # Tool Management # ====================================================================================================================== @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) @raise_on_invalid_id(param_name="tool_id", expected_prefix=PrimitiveType.TOOL) + @trace_method async def attach_tool_async(self, agent_id: str, tool_id: str, actor: PydanticUser) -> None: """ Attaches a tool to an agent. @@ -2573,8 +2573,8 @@ class AgentManager: await session.commit() @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def bulk_attach_tools_async(self, agent_id: str, tool_ids: List[str], actor: PydanticUser) -> None: """ Efficiently attaches multiple tools to an agent in a single operation. @@ -2739,9 +2739,9 @@ class AgentManager: return PydanticAgentState(**agent_state_dict) @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) @raise_on_invalid_id(param_name="tool_id", expected_prefix=PrimitiveType.TOOL) + @trace_method async def detach_tool_async(self, agent_id: str, tool_id: str, actor: PydanticUser) -> None: """ Detaches a tool from an agent. @@ -2770,8 +2770,8 @@ class AgentManager: await session.commit() @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def bulk_detach_tools_async(self, agent_id: str, tool_ids: List[str], actor: PydanticUser) -> None: """ Efficiently detaches multiple tools from an agent in a single operation. @@ -2807,8 +2807,8 @@ class AgentManager: await session.commit() @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def modify_approvals_async(self, agent_id: str, tool_name: str, requires_approval: bool, actor: PydanticUser) -> None: def is_target_rule(rule): return rule.tool_name == tool_name and rule.type == "requires_approval" @@ -3157,8 +3157,8 @@ class AgentManager: return results @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def get_agent_files_config_async(self, agent_id: str, actor: PydanticUser) -> Tuple[int, int]: """Get per_file_view_window_char_limit and max_files_open for an agent. @@ -3214,8 +3214,8 @@ class AgentManager: return per_file_limit, max_files @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def get_agent_max_files_open_async(self, agent_id: str, actor: PydanticUser) -> int: """Get max_files_open for an agent. @@ -3243,8 +3243,8 @@ class AgentManager: return row @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def get_agent_per_file_view_window_char_limit_async(self, agent_id: str, actor: PydanticUser) -> int: """Get per_file_view_window_char_limit for an agent. @@ -3272,8 +3272,8 @@ class AgentManager: return row @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def get_context_window(self, agent_id: str, actor: PydanticUser) -> ContextWindowOverview: agent_state, system_message, num_messages, num_archival_memories = await self.rebuild_system_prompt_async( agent_id=agent_id, actor=actor, force=True, dry_run=True diff --git a/letta/services/archive_manager.py b/letta/services/archive_manager.py index 5ed6220f..79c09c8e 100644 --- a/letta/services/archive_manager.py +++ b/letta/services/archive_manager.py @@ -55,8 +55,8 @@ class ArchiveManager: raise @enforce_types - @trace_method @raise_on_invalid_id(param_name="archive_id", expected_prefix=PrimitiveType.ARCHIVE) + @trace_method async def get_archive_by_id_async( self, archive_id: str, @@ -72,8 +72,8 @@ class ArchiveManager: return archive.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="archive_id", expected_prefix=PrimitiveType.ARCHIVE) + @trace_method async def update_archive_async( self, archive_id: str, @@ -99,8 +99,8 @@ class ArchiveManager: return archive.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def list_archives_async( self, *, @@ -150,9 +150,9 @@ class ArchiveManager: return [a.to_pydantic() for a in archives] @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) @raise_on_invalid_id(param_name="archive_id", expected_prefix=PrimitiveType.ARCHIVE) + @trace_method async def attach_agent_to_archive_async( self, agent_id: str, @@ -194,9 +194,9 @@ class ArchiveManager: await session.commit() @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) @raise_on_invalid_id(param_name="archive_id", expected_prefix=PrimitiveType.ARCHIVE) + @trace_method async def detach_agent_from_archive_async( self, agent_id: str, @@ -227,8 +227,8 @@ class ArchiveManager: await session.commit() @enforce_types - @trace_method @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def get_default_archive_for_agent_async( self, agent_id: str, @@ -260,8 +260,8 @@ class ArchiveManager: return None @enforce_types - @trace_method @raise_on_invalid_id(param_name="archive_id", expected_prefix=PrimitiveType.ARCHIVE) + @trace_method async def delete_archive_async( self, archive_id: str, @@ -278,8 +278,8 @@ class ArchiveManager: logger.info(f"Deleted archive {archive_id}") @enforce_types - @trace_method @raise_on_invalid_id(param_name="archive_id", expected_prefix=PrimitiveType.ARCHIVE) + @trace_method async def create_passage_in_archive_async( self, archive_id: str, @@ -360,9 +360,9 @@ class ArchiveManager: return created_passage @enforce_types - @trace_method @raise_on_invalid_id(param_name="archive_id", expected_prefix=PrimitiveType.ARCHIVE) @raise_on_invalid_id(param_name="passage_id", expected_prefix=PrimitiveType.PASSAGE) + @trace_method async def delete_passage_from_archive_async( self, archive_id: str, @@ -470,8 +470,8 @@ class ArchiveManager: raise @enforce_types - @trace_method @raise_on_invalid_id(param_name="archive_id", expected_prefix=PrimitiveType.ARCHIVE) + @trace_method async def get_agents_for_archive_async( self, archive_id: str, @@ -583,8 +583,8 @@ class ArchiveManager: return agent_ids[0] @enforce_types - @trace_method @raise_on_invalid_id(param_name="archive_id", expected_prefix=PrimitiveType.ARCHIVE) + @trace_method async def get_or_set_vector_db_namespace_async( self, archive_id: str, diff --git a/letta/services/block_manager.py b/letta/services/block_manager.py index 7014dd7e..1ed2392f 100644 --- a/letta/services/block_manager.py +++ b/letta/services/block_manager.py @@ -134,8 +134,8 @@ class BlockManager: return result @enforce_types - @trace_method @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK) + @trace_method async def update_block_async(self, block_id: str, block_update: BlockUpdate, actor: PydanticUser) -> PydanticBlock: """Update a block by its ID with the given BlockUpdate object.""" async with db_registry.async_session() as session: @@ -154,8 +154,8 @@ class BlockManager: return pydantic_block @enforce_types - @trace_method @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK) + @trace_method async def delete_block_async(self, block_id: str, actor: PydanticUser) -> None: """Delete a block by its ID.""" async with db_registry.async_session() as session: @@ -353,8 +353,8 @@ class BlockManager: return [block.to_pydantic() for block in blocks] @enforce_types - @trace_method @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK) + @trace_method async def get_block_by_id_async(self, block_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticBlock]: """Retrieve a block by its name.""" async with db_registry.async_session() as session: @@ -413,8 +413,8 @@ class BlockManager: return pydantic_blocks @enforce_types - @trace_method @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK) + @trace_method async def get_agents_for_block_async( self, block_id: str, @@ -600,9 +600,9 @@ class BlockManager: return None @enforce_types - @trace_method @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK) @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def checkpoint_block_async( self, block_id: str, @@ -710,8 +710,8 @@ class BlockManager: return updated_block @enforce_types - @trace_method @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK) + @trace_method async def undo_checkpoint_block( self, block_id: str, actor: PydanticUser, use_preloaded_block: Optional[BlockModel] = None ) -> PydanticBlock: @@ -761,8 +761,8 @@ class BlockManager: return block.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK) + @trace_method async def redo_checkpoint_block( self, block_id: str, actor: PydanticUser, use_preloaded_block: Optional[BlockModel] = None ) -> PydanticBlock: diff --git a/letta/services/file_manager.py b/letta/services/file_manager.py index d1d14d13..6301927a 100644 --- a/letta/services/file_manager.py +++ b/letta/services/file_manager.py @@ -93,8 +93,8 @@ class FileManager: # 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="file_id", expected_prefix=PrimitiveType.FILE) + @trace_method # @async_redis_cache( # key_func=lambda self, file_id, actor=None, include_content=False, strip_directory_prefix=False: f"{file_id}:{actor.organization_id if actor else 'none'}:{include_content}:{strip_directory_prefix}", # prefix="file_content", @@ -136,8 +136,8 @@ class FileManager: return await file_orm.to_pydantic_async(include_content=include_content, strip_directory_prefix=strip_directory_prefix) @enforce_types - @trace_method @raise_on_invalid_id(param_name="file_id", expected_prefix=PrimitiveType.FILE) + @trace_method async def update_file_status( self, *, @@ -354,8 +354,8 @@ class FileManager: return file_metadata @enforce_types - @trace_method @raise_on_invalid_id(param_name="file_id", expected_prefix=PrimitiveType.FILE) + @trace_method async def upsert_file_content( self, *, @@ -400,8 +400,8 @@ class FileManager: return await result.scalar_one().to_pydantic_async(include_content=True) @enforce_types - @trace_method @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE) + @trace_method async def list_files( self, source_id: str, @@ -462,8 +462,8 @@ class FileManager: return file_metadatas @enforce_types - @trace_method @raise_on_invalid_id(param_name="file_id", expected_prefix=PrimitiveType.FILE) + @trace_method async def delete_file(self, file_id: str, actor: PydanticUser) -> PydanticFileMetadata: """Delete a file by its ID.""" async with db_registry.async_session() as session: @@ -517,8 +517,8 @@ class FileManager: return f"{source.name}/{base}_({count}){ext}" @enforce_types - @trace_method @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE) + @trace_method # @async_redis_cache( # key_func=lambda self, original_filename, source_id, actor: f"{original_filename}:{source_id}:{actor.organization_id}", # prefix="file_by_name", diff --git a/letta/services/group_manager.py b/letta/services/group_manager.py index 4dbdb149..6074302f 100644 --- a/letta/services/group_manager.py +++ b/letta/services/group_manager.py @@ -65,8 +65,8 @@ class GroupManager: return [group.to_pydantic() for group in groups] @enforce_types - @trace_method @raise_on_invalid_id(param_name="group_id", expected_prefix=PrimitiveType.GROUP) + @trace_method async def retrieve_group_async(self, group_id: str, actor: PydanticUser) -> PydanticGroup: async with db_registry.async_session() as session: group = await GroupModel.read_async(db_session=session, identifier=group_id, actor=actor) @@ -123,8 +123,8 @@ class GroupManager: return new_group.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="group_id", expected_prefix=PrimitiveType.GROUP) + @trace_method async def modify_group_async(self, group_id: str, group_update: GroupUpdate, actor: PydanticUser) -> PydanticGroup: async with db_registry.async_session() as session: group = await GroupModel.read_async(db_session=session, identifier=group_id, actor=actor) @@ -187,16 +187,16 @@ class GroupManager: return group.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="group_id", expected_prefix=PrimitiveType.GROUP) + @trace_method async def delete_group_async(self, group_id: str, actor: PydanticUser) -> None: async with db_registry.async_session() as session: group = await GroupModel.read_async(db_session=session, identifier=group_id, actor=actor) await group.hard_delete_async(session) @enforce_types - @trace_method @raise_on_invalid_id(param_name="group_id", expected_prefix=PrimitiveType.GROUP) + @trace_method async def list_group_messages_async( self, actor: PydanticUser, @@ -233,8 +233,8 @@ class GroupManager: return messages @enforce_types - @trace_method @raise_on_invalid_id(param_name="group_id", expected_prefix=PrimitiveType.GROUP) + @trace_method async def reset_messages_async(self, group_id: str, actor: PydanticUser) -> None: async with db_registry.async_session() as session: # Ensure group is loadable by user @@ -249,8 +249,8 @@ class GroupManager: await session.commit() @enforce_types - @trace_method @raise_on_invalid_id(param_name="group_id", expected_prefix=PrimitiveType.GROUP) + @trace_method async def bump_turns_counter_async(self, group_id: str, actor: PydanticUser) -> int: async with db_registry.async_session() as session: # Ensure group is loadable by user @@ -262,9 +262,9 @@ class GroupManager: return group.turns_counter @enforce_types - @trace_method @raise_on_invalid_id(param_name="group_id", expected_prefix=PrimitiveType.GROUP) @raise_on_invalid_id(param_name="last_processed_message_id", expected_prefix=PrimitiveType.MESSAGE) + @trace_method async def get_last_processed_message_id_and_update_async( self, group_id: str, last_processed_message_id: str, actor: PydanticUser ) -> str: @@ -413,9 +413,9 @@ class GroupManager: session.add(BlocksAgents(agent_id=manager_agent.id, block_id=block.id, block_label=block.label)) @enforce_types - @trace_method @raise_on_invalid_id(param_name="group_id", expected_prefix=PrimitiveType.GROUP) @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK) + @trace_method async def attach_block_async(self, group_id: str, block_id: str, actor: PydanticUser) -> None: """Attach a block to a group.""" async with db_registry.async_session() as session: @@ -437,9 +437,9 @@ class GroupManager: await session.commit() @enforce_types - @trace_method @raise_on_invalid_id(param_name="group_id", expected_prefix=PrimitiveType.GROUP) @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK) + @trace_method async def detach_block_async(self, group_id: str, block_id: str, actor: PydanticUser) -> None: """Detach a block from a group.""" async with db_registry.async_session() as session: diff --git a/letta/services/identity_manager.py b/letta/services/identity_manager.py index 636d8de5..91adebea 100644 --- a/letta/services/identity_manager.py +++ b/letta/services/identity_manager.py @@ -83,8 +83,8 @@ class IdentityManager: return [identity.to_pydantic() for identity in identities], next_cursor, has_more @enforce_types - @trace_method @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY) + @trace_method async def get_identity_async(self, identity_id: str, actor: PydanticUser) -> PydanticIdentity: async with db_registry.async_session() as session: identity = await IdentityModel.read_async(db_session=session, identifier=identity_id, actor=actor) @@ -165,8 +165,8 @@ class IdentityManager: ) @enforce_types - @trace_method @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY) + @trace_method async def update_identity_async( self, identity_id: str, identity: IdentityUpdate, actor: PydanticUser, replace: bool = False ) -> PydanticIdentity: @@ -229,8 +229,8 @@ class IdentityManager: return existing_identity.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY) + @trace_method async def upsert_identity_properties_async( self, identity_id: str, properties: List[IdentityProperty], actor: PydanticUser ) -> PydanticIdentity: @@ -247,8 +247,8 @@ class IdentityManager: ) @enforce_types - @trace_method @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY) + @trace_method async def delete_identity_async(self, identity_id: str, actor: PydanticUser) -> None: async with db_registry.async_session() as session: identity = await IdentityModel.read_async(db_session=session, identifier=identity_id, actor=actor) @@ -305,8 +305,8 @@ class IdentityManager: current_relationship.extend(new_items) @enforce_types - @trace_method @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY) + @trace_method async def list_agents_for_identity_async( self, identity_id: str, @@ -338,8 +338,8 @@ class IdentityManager: return await asyncio.gather(*[agent.to_pydantic_async(include_relationships=[], include=include) for agent in agents]) @enforce_types - @trace_method @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY) + @trace_method async def list_blocks_for_identity_async( self, identity_id: str, @@ -370,9 +370,9 @@ class IdentityManager: return [block.to_pydantic() for block in blocks] @enforce_types - @trace_method @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY) @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def attach_agent_async(self, identity_id: str, agent_id: str, actor: PydanticUser) -> None: """ Attach an agent to an identity. @@ -388,9 +388,9 @@ class IdentityManager: await identity.update_async(db_session=session, actor=actor) @enforce_types - @trace_method @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY) @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT) + @trace_method async def detach_agent_async(self, identity_id: str, agent_id: str, actor: PydanticUser) -> None: """ Detach an agent from an identity. @@ -406,9 +406,9 @@ class IdentityManager: await identity.update_async(db_session=session, actor=actor) @enforce_types - @trace_method @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY) @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK) + @trace_method async def attach_block_async(self, identity_id: str, block_id: str, actor: PydanticUser) -> None: """ Attach a block to an identity. @@ -424,9 +424,9 @@ class IdentityManager: await identity.update_async(db_session=session, actor=actor) @enforce_types - @trace_method @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY) @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK) + @trace_method async def detach_block_async(self, identity_id: str, block_id: str, actor: PydanticUser) -> None: """ Detach a block from an identity. diff --git a/letta/services/job_manager.py b/letta/services/job_manager.py index 33a6e674..ddabc80f 100644 --- a/letta/services/job_manager.py +++ b/letta/services/job_manager.py @@ -70,8 +70,8 @@ class JobManager: return result @enforce_types - @trace_method @raise_on_invalid_id(param_name="job_id", expected_prefix=PrimitiveType.JOB) + @trace_method async def update_job_by_id_async( self, job_id: str, job_update: JobUpdate, actor: PydanticUser, safe_update: bool = False ) -> PydanticJob: @@ -148,8 +148,8 @@ class JobManager: return result @enforce_types - @trace_method @raise_on_invalid_id(param_name="job_id", expected_prefix=PrimitiveType.JOB) + @trace_method async def safe_update_job_status_async( self, job_id: str, @@ -189,8 +189,8 @@ class JobManager: return False @enforce_types - @trace_method @raise_on_invalid_id(param_name="job_id", expected_prefix=PrimitiveType.JOB) + @trace_method async def get_job_by_id_async(self, job_id: str, actor: PydanticUser) -> PydanticJob: """Fetch a job by its ID asynchronously.""" async with db_registry.async_session() as session: @@ -304,8 +304,8 @@ class JobManager: return [job.to_pydantic() for job in jobs] @enforce_types - @trace_method @raise_on_invalid_id(param_name="job_id", expected_prefix=PrimitiveType.JOB) + @trace_method async def delete_job_by_id_async(self, job_id: str, actor: PydanticUser) -> PydanticJob: """Delete a job by its ID.""" async with db_registry.async_session() as session: @@ -314,8 +314,8 @@ class JobManager: return job.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="run_id", expected_prefix=PrimitiveType.RUN) + @trace_method async def get_run_messages( self, run_id: str, @@ -372,8 +372,8 @@ class JobManager: return messages @enforce_types - @trace_method @raise_on_invalid_id(param_name="run_id", expected_prefix=PrimitiveType.RUN) + @trace_method async def get_step_messages( self, run_id: str, @@ -537,8 +537,8 @@ class JobManager: return result @enforce_types - @trace_method @raise_on_invalid_id(param_name="job_id", expected_prefix=PrimitiveType.JOB) + @trace_method async def get_job_steps( self, job_id: str, diff --git a/letta/services/message_manager.py b/letta/services/message_manager.py index 2021ca61..f0ee955f 100644 --- a/letta/services/message_manager.py +++ b/letta/services/message_manager.py @@ -345,8 +345,8 @@ class MessageManager: return combined_messages @enforce_types - @trace_method @raise_on_invalid_id(param_name="message_id", expected_prefix=PrimitiveType.MESSAGE) + @trace_method async def get_message_by_id_async(self, message_id: str, actor: PydanticUser) -> Optional[PydanticMessage]: """Fetch a message by ID.""" async with db_registry.async_session() as session: @@ -754,8 +754,8 @@ class MessageManager: return message @enforce_types - @trace_method @raise_on_invalid_id(param_name="message_id", expected_prefix=PrimitiveType.MESSAGE) + @trace_method async def delete_message_by_id_async(self, message_id: str, actor: PydanticUser, strict_mode: bool = False) -> bool: """Delete a message (async version with turbopuffer support).""" # capture agent_id before deletion diff --git a/letta/services/provider_manager.py b/letta/services/provider_manager.py index 50d95201..8fbaf043 100644 --- a/letta/services/provider_manager.py +++ b/letta/services/provider_manager.py @@ -95,8 +95,8 @@ class ProviderManager: return provider_pydantic @enforce_types - @trace_method @raise_on_invalid_id(param_name="provider_id", expected_prefix=PrimitiveType.PROVIDER) + @trace_method async def update_provider_async(self, provider_id: str, provider_update: ProviderUpdate, actor: PydanticUser) -> PydanticProvider: """Update provider details.""" async with db_registry.async_session() as session: @@ -151,8 +151,8 @@ class ProviderManager: return existing_provider.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="provider_id", expected_prefix=PrimitiveType.PROVIDER) + @trace_method async def delete_provider_by_id_async(self, provider_id: str, actor: PydanticUser): """Delete a provider.""" async with db_registry.async_session() as session: @@ -284,8 +284,8 @@ class ProviderManager: return [provider.to_pydantic() for provider in all_providers] @enforce_types - @trace_method @raise_on_invalid_id(param_name="provider_id", expected_prefix=PrimitiveType.PROVIDER) + @trace_method async def get_provider_async(self, provider_id: str, actor: PydanticUser) -> PydanticProvider: async with db_registry.async_session() as session: # First try to get as organization-specific provider diff --git a/letta/services/sandbox_config_manager.py b/letta/services/sandbox_config_manager.py index 3ff68845..cf26608e 100644 --- a/letta/services/sandbox_config_manager.py +++ b/letta/services/sandbox_config_manager.py @@ -101,8 +101,8 @@ class SandboxConfigManager: return db_sandbox.to_pydantic() @enforce_types - @trace_method @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: @@ -130,8 +130,8 @@ class SandboxConfigManager: return sandbox.to_pydantic() @enforce_types - @trace_method @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: @@ -178,8 +178,8 @@ class SandboxConfigManager: return None @enforce_types - @trace_method @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: @@ -267,8 +267,8 @@ class SandboxConfigManager: return await PydanticEnvVar.from_orm_async(env_var) @enforce_types - @trace_method @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, @@ -310,8 +310,8 @@ class SandboxConfigManager: return result @enforce_types - @trace_method @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]: @@ -323,8 +323,8 @@ class SandboxConfigManager: return result @enforce_types - @trace_method @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]: @@ -333,8 +333,8 @@ class SandboxConfigManager: return {env_var.key: env_var.value for env_var in env_vars} @enforce_types - @trace_method @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: Optional[PydanticUser] = None ) -> Optional[PydanticEnvVar]: diff --git a/letta/services/source_manager.py b/letta/services/source_manager.py index d79e3c1c..0a3b7b3a 100644 --- a/letta/services/source_manager.py +++ b/letta/services/source_manager.py @@ -201,8 +201,8 @@ class SourceManager: return sources @enforce_types - @trace_method @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE) + @trace_method 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: @@ -225,8 +225,8 @@ class SourceManager: return source.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE) + @trace_method 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: @@ -270,8 +270,8 @@ class SourceManager: return await SourceModel.size_async(db_session=session, actor=actor) @enforce_types - @trace_method @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE) + @trace_method async def list_attached_agents( self, source_id: str, actor: PydanticUser, ids_only: bool = False ) -> Union[List[PydanticAgentState], List[str]]: @@ -324,8 +324,8 @@ class SourceManager: return await asyncio.gather(*[agent.to_pydantic_async() for agent in agents_orm]) @enforce_types - @trace_method @raise_on_invalid_id(param_name="source_id", expected_prefix=PrimitiveType.SOURCE) + @trace_method async def get_agents_for_source_id( self, source_id: str, @@ -439,8 +439,8 @@ 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) + @trace_method 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: diff --git a/letta/services/step_manager.py b/letta/services/step_manager.py index 605717eb..cd57bf5f 100644 --- a/letta/services/step_manager.py +++ b/letta/services/step_manager.py @@ -33,9 +33,9 @@ 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) + @trace_method async def list_steps_async( self, actor: PydanticUser, @@ -82,11 +82,11 @@ class StepManager: return [step.to_pydantic() for step in steps] @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) + @trace_method def log_step( self, actor: PydanticUser, @@ -140,11 +140,11 @@ class StepManager: return new_step.to_pydantic() @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) + @trace_method async def log_step_async( self, actor: PydanticUser, @@ -207,24 +207,24 @@ class StepManager: return pydantic_step @enforce_types - @trace_method @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP) + @trace_method 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) return step.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP) + @trace_method 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) return metrics.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP) + @trace_method async def add_feedback_async( self, step_id: str, feedback: FeedbackType | None, actor: PydanticUser, tags: list[str] | None = None ) -> PydanticStep: @@ -239,8 +239,8 @@ class StepManager: return step.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP) + @trace_method async def update_step_transaction_id(self, actor: PydanticUser, step_id: str, transaction_id: str) -> PydanticStep: """Update the transaction ID for a step. @@ -267,8 +267,8 @@ class StepManager: return step.to_pydantic() @enforce_types - @trace_method @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP) + @trace_method async def list_step_messages_async( self, step_id: str, @@ -291,8 +291,8 @@ class StepManager: return [message.to_pydantic() for message in messages] @enforce_types - @trace_method @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP) + @trace_method async def update_step_stop_reason(self, actor: PydanticUser, step_id: str, stop_reason: StopReasonType) -> PydanticStep: """Update the stop reason for a step. @@ -319,8 +319,8 @@ class StepManager: return step @enforce_types - @trace_method @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP) + @trace_method async def update_step_error_async( self, actor: PydanticUser, @@ -369,8 +369,8 @@ class StepManager: return pydantic_step @enforce_types - @trace_method @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP) + @trace_method async def update_step_success_async( self, actor: PydanticUser, @@ -420,8 +420,8 @@ class StepManager: return pydantic_step @enforce_types - @trace_method @raise_on_invalid_id(param_name="step_id", expected_prefix=PrimitiveType.STEP) + @trace_method async def update_step_cancelled_async( self, actor: PydanticUser, @@ -460,10 +460,10 @@ class StepManager: return pydantic_step @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) + @trace_method async def record_step_metrics_async( self, actor: PydanticUser,