From 3abb1d05807753f950af655798307c6fe19ee1c6 Mon Sep 17 00:00:00 2001 From: Matthew Zhou Date: Tue, 11 Mar 2025 14:37:17 -0700 Subject: [PATCH] fix: BlockUpdate resets limit due to hardcoding (#1246) --- letta/schemas/block.py | 52 ++-------------------- letta/server/rest_api/routers/v1/agents.py | 2 +- tests/test_managers.py | 32 +++++++++---- 3 files changed, 28 insertions(+), 58 deletions(-) diff --git a/letta/schemas/block.py b/letta/schemas/block.py index 2aa518cb..3e2fbb7e 100644 --- a/letta/schemas/block.py +++ b/letta/schemas/block.py @@ -1,6 +1,6 @@ from typing import Optional -from pydantic import BaseModel, Field, model_validator +from pydantic import Field, model_validator from typing_extensions import Self from letta.constants import CORE_MEMORY_BLOCK_CHAR_LIMIT @@ -37,7 +37,8 @@ class BaseBlock(LettaBase, validate_assignment=True): @model_validator(mode="after") def verify_char_limit(self) -> Self: - if self.value and len(self.value) > self.limit: + # self.limit can be None from + if self.limit is not None and self.value and len(self.value) > self.limit: error_msg = f"Edit failed: Exceeds {self.limit} character limit (requested {len(self.value)}) - {str(self)}." raise ValueError(error_msg) @@ -89,61 +90,16 @@ class Persona(Block): label: str = "persona" -# class CreateBlock(BaseBlock): -# """Create a block""" -# -# is_template: bool = True -# label: str = Field(..., description="Label of the block.") - - -class BlockLabelUpdate(BaseModel): - """Update the label of a block""" - - current_label: str = Field(..., description="Current label of the block.") - new_label: str = Field(..., description="New label of the block.") - - -# class CreatePersona(CreateBlock): -# """Create a persona block""" -# -# label: str = "persona" -# -# -# class CreateHuman(CreateBlock): -# """Create a human block""" -# -# label: str = "human" - - class BlockUpdate(BaseBlock): """Update a block""" - limit: Optional[int] = Field(CORE_MEMORY_BLOCK_CHAR_LIMIT, description="Character limit of the block.") + limit: Optional[int] = Field(None, description="Character limit of the block.") value: Optional[str] = Field(None, description="Value of the block.") class Config: extra = "ignore" # Ignores extra fields -class BlockLimitUpdate(BaseModel): - """Update the limit of a block""" - - label: str = Field(..., description="Label of the block.") - limit: int = Field(..., description="New limit of the block.") - - -# class UpdatePersona(BlockUpdate): -# """Update a persona block""" -# -# label: str = "persona" -# -# -# class UpdateHuman(BlockUpdate): -# """Update a human block""" -# -# label: str = "human" - - class CreateBlock(BaseBlock): """Create a block""" diff --git a/letta/server/rest_api/routers/v1/agents.py b/letta/server/rest_api/routers/v1/agents.py index 2b7aec1c..5de351a2 100644 --- a/letta/server/rest_api/routers/v1/agents.py +++ b/letta/server/rest_api/routers/v1/agents.py @@ -13,7 +13,7 @@ from letta.constants import DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG from letta.log import get_logger from letta.orm.errors import NoResultFound from letta.schemas.agent import AgentState, CreateAgent, UpdateAgent -from letta.schemas.block import Block, BlockUpdate, CreateBlock # , BlockLabelUpdate, BlockLimitUpdate +from letta.schemas.block import Block, BlockUpdate from letta.schemas.job import JobStatus, JobUpdate, LettaRequestConfig from letta.schemas.letta_message import LettaMessageUnion, LettaMessageUpdateUnion from letta.schemas.letta_request import LettaRequest, LettaStreamingRequest diff --git a/tests/test_managers.py b/tests/test_managers.py index 61bfad92..49744d62 100644 --- a/tests/test_managers.py +++ b/tests/test_managers.py @@ -2149,28 +2149,42 @@ def test_update_block(server: SyncServer, default_user): def test_update_block_limit(server: SyncServer, default_user): - block_manager = BlockManager() block = block_manager.create_or_update_block(PydanticBlock(label="persona", value="Original Content"), actor=default_user) limit = len("Updated Content") * 2000 - update_data = BlockUpdate(value="Updated Content" * 2000, description="Updated description", limit=limit) + update_data = BlockUpdate(value="Updated Content" * 2000, description="Updated description") - # Check that a large block fails - try: + # Check that exceeding the block limit raises an exception + with pytest.raises(ValueError): block_manager.update_block(block_id=block.id, block_update=update_data, actor=default_user) - assert False - except Exception: - pass + # Ensure the update works when within limits + update_data = BlockUpdate(value="Updated Content" * 2000, description="Updated description", limit=limit) block_manager.update_block(block_id=block.id, block_update=update_data, actor=default_user) - # Retrieve the updated block + + # Retrieve the updated block and validate the update updated_block = block_manager.get_blocks(actor=default_user, id=block.id)[0] - # Assertions to verify the update + assert updated_block.value == "Updated Content" * 2000 assert updated_block.description == "Updated description" +def test_update_block_limit_does_not_reset(server: SyncServer, default_user): + block_manager = BlockManager() + new_content = "Updated Content" * 2000 + limit = len(new_content) + block = block_manager.create_or_update_block(PydanticBlock(label="persona", value="Original Content", limit=limit), actor=default_user) + + # Ensure the update works + update_data = BlockUpdate(value=new_content) + block_manager.update_block(block_id=block.id, block_update=update_data, actor=default_user) + + # Retrieve the updated block and validate the update + updated_block = block_manager.get_blocks(actor=default_user, id=block.id)[0] + assert updated_block.value == new_content + + def test_delete_block(server: SyncServer, default_user): block_manager = BlockManager()