Fimeg/letta-server
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix: correct decorator order for trace_method and raise_on_invalid_id (#7226)

Browse Source 
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)
This commit is contained in:
jnjpng
2025-12-16 16:10:24 -08:00
committed by Caren Thomas
parent 82de08be42
commit 591420876a
12 changed files with 102 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
letta/services/agent_manager.py
View File

@@ -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

22
letta/services/archive_manager.py
View File

@@ -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,

14
letta/services/block_manager.py
View File

@@ -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:

12
letta/services/file_manager.py
View File

@@ -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",

18
letta/services/group_manager.py
View File

@@ -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:

20
letta/services/identity_manager.py
View File

@@ -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.

14
letta/services/job_manager.py
View File

@@ -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,

4
letta/services/message_manager.py
View File

@@ -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

6
letta/services/provider_manager.py
View File

@@ -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

14
letta/services/sandbox_config_manager.py
View File

@@ -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]:

10
letta/services/source_manager.py
View File

@@ -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:

26
letta/services/step_manager.py
View File

@@ -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,