From 7346473ebe0e470f2e22cf72ff1b71f737e4413a Mon Sep 17 00:00:00 2001 From: cthomas Date: Wed, 28 May 2025 23:15:36 -0700 Subject: [PATCH] feat: optimize get message ids (#2508) --- letta/services/message_manager.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/letta/services/message_manager.py b/letta/services/message_manager.py index 26e5326c..4141a4bc 100644 --- a/letta/services/message_manager.py +++ b/letta/services/message_manager.py @@ -49,7 +49,7 @@ class MessageManager: def get_messages_by_ids(self, message_ids: List[str], actor: PydanticUser) -> List[PydanticMessage]: """Fetch messages by ID and return them in the requested order.""" with db_registry.session() as session: - results = MessageModel.list(db_session=session, id=message_ids, organization_id=actor.organization_id, limit=len(message_ids)) + results = MessageModel.read_multiple(db_session=session, identifiers=message_ids, actor=actor) return self._get_messages_by_id_postprocess(results, message_ids) @enforce_types @@ -57,9 +57,7 @@ class MessageManager: async def get_messages_by_ids_async(self, message_ids: List[str], actor: PydanticUser) -> List[PydanticMessage]: """Fetch messages by ID and return them in the requested order. Async version of above function.""" async with db_registry.async_session() as session: - results = await MessageModel.list_async( - db_session=session, id=message_ids, organization_id=actor.organization_id, limit=len(message_ids) - ) + results = await MessageModel.read_multiple_async(db_session=session, identifiers=message_ids, actor=actor) return self._get_messages_by_id_postprocess(results, message_ids) def _get_messages_by_id_postprocess(