From edbddd746dc43e5eaecf80662ba6bed91911ea2e Mon Sep 17 00:00:00 2001 From: cthomas Date: Fri, 23 May 2025 09:53:43 -0700 Subject: [PATCH] feat(asyncify): migrate count blocks (#2374) --- letta/server/rest_api/routers/v1/blocks.py | 5 +++-- letta/services/block_manager.py | 9 +++------ 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/letta/server/rest_api/routers/v1/blocks.py b/letta/server/rest_api/routers/v1/blocks.py index d31fd855..b09097ad 100644 --- a/letta/server/rest_api/routers/v1/blocks.py +++ b/letta/server/rest_api/routers/v1/blocks.py @@ -39,14 +39,15 @@ async def list_blocks( @router.get("/count", response_model=int, operation_id="count_blocks") -def count_blocks( +async def count_blocks( server: SyncServer = Depends(get_letta_server), actor_id: Optional[str] = Header(None, alias="user_id"), ): """ Count all blocks created by a user. """ - return server.block_manager.size(actor=server.user_manager.get_user_or_default(user_id=actor_id)) + actor = await server.user_manager.get_actor_or_default_async(actor_id=actor_id) + return await server.block_manager.size_async(actor=actor) @router.post("/", response_model=Block, operation_id="create_block") diff --git a/letta/services/block_manager.py b/letta/services/block_manager.py index fd46e86a..3a642ba1 100644 --- a/letta/services/block_manager.py +++ b/letta/services/block_manager.py @@ -273,15 +273,12 @@ class BlockManager: @trace_method @enforce_types - def size( - self, - actor: PydanticUser, - ) -> int: + async def size_async(self, actor: PydanticUser) -> int: """ Get the total count of blocks for the given user. """ - with db_registry.session() as session: - return BlockModel.size(db_session=session, actor=actor) + async with db_registry.async_session() as session: + return await BlockModel.size_async(db_session=session, actor=actor) # Block History Functions