From f47d7e557dbd4b4bf0360d0a9ef376d76b8d595c Mon Sep 17 00:00:00 2001 From: cthomas Date: Thu, 29 May 2025 09:39:32 -0700 Subject: [PATCH] feat: make size load concurrent (#2520) --- letta/agents/letta_agent.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/letta/agents/letta_agent.py b/letta/agents/letta_agent.py index 09981168..59693a20 100644 --- a/letta/agents/letta_agent.py +++ b/letta/agents/letta_agent.py @@ -473,9 +473,17 @@ class LettaAgent(BaseAgent): agent_state: AgentState, tool_rules_solver: ToolRulesSolver, ) -> ChatCompletion | AsyncStream[ChatCompletionChunk]: - self.num_messages = self.num_messages or (await self.message_manager.size_async(actor=self.actor, agent_id=agent_state.id)) - self.num_archival_memories = self.num_archival_memories or ( - await self.passage_manager.size_async(actor=self.actor, agent_id=agent_state.id) + self.num_messages, self.num_archival_memories = await asyncio.gather( + ( + self.message_manager.size_async(actor=self.actor, agent_id=agent_state.id) + if self.num_messages is None + else asyncio.sleep(0, result=self.num_messages) + ), + ( + self.passage_manager.size_async(actor=self.actor, agent_id=agent_state.id) + if self.num_archival_memories is None + else asyncio.sleep(0, result=self.num_archival_memories) + ), ) in_context_messages = await self._rebuild_memory_async( in_context_messages, agent_state, num_messages=self.num_messages, num_archival_memories=self.num_archival_memories