feat (asyncify): async batch block creation for agent creation (#2397)

This commit is contained in:
Sarah Wooders
2025-05-23 18:07:04 -07:00
committed by GitHub
parent 28c3624a88
commit 06ecf8ede9
2 changed files with 25 additions and 1 deletions

View File

@@ -387,7 +387,7 @@ class AgentManager:
block_ids = list(agent_create.block_ids or [])
if agent_create.memory_blocks:
pydantic_blocks = [PydanticBlock(**b.model_dump(to_orm=True)) for b in agent_create.memory_blocks]
created_blocks = self.block_manager.batch_create_blocks(
created_blocks = await self.block_manager.batch_create_blocks_async(
pydantic_blocks,
actor=actor,
)

View File

@@ -77,6 +77,30 @@ class BlockManager:
# Convert back to Pydantic
return [m.to_pydantic() for m in created_models]
@trace_method
@enforce_types
async def batch_create_blocks_async(self, blocks: List[PydanticBlock], actor: PydanticUser) -> List[PydanticBlock]:
"""
Batch-create multiple Blocks in one transaction for better performance.
Args:
blocks: List of PydanticBlock schemas to create
actor: The user performing the operation
Returns:
List of created PydanticBlock instances (with IDs, timestamps, etc.)
"""
if not blocks:
return []
async with db_registry.async_session() as session:
block_models = [
BlockModel(**block.model_dump(to_orm=True, exclude_none=True), organization_id=actor.organization_id) for block in blocks
]
created_models = await BlockModel.batch_create_async(items=block_models, db_session=session, actor=actor)
# Convert back to Pydantic
return [m.to_pydantic() for m in created_models]
@trace_method
@enforce_types
def update_block(self, block_id: str, block_update: BlockUpdate, actor: PydanticUser) -> PydanticBlock: