From e62ddf8355a25751d31164628b575231c8e9635f Mon Sep 17 00:00:00 2001 From: Matthew Zhou Date: Tue, 22 Jul 2025 10:43:37 -0700 Subject: [PATCH] feat: Add per-agent file management controls with context-aware defaults (#3467) --- ...907b38_add_file_controls_to_agent_state.py | 33 ++ letta/constants.py | 6 +- letta/orm/agent.py | 35 ++- letta/orm/files_agents.py | 10 +- letta/schemas/agent.py | 49 ++- letta/schemas/agent_file.py | 2 + letta/server/server.py | 2 + letta/services/agent_file_manager.py | 18 +- letta/services/agent_manager.py | 79 ++++- letta/services/files_agents_manager.py | 35 ++- .../services/helpers/agent_manager_helper.py | 9 +- .../tool_executor/files_tool_executor.py | 17 +- letta/utils.py | 35 ++- tests/test_managers.py | 292 ++++++++++++++---- tests/test_sources.py | 4 +- 15 files changed, 526 insertions(+), 100 deletions(-) create mode 100644 alembic/versions/c4eb5a907b38_add_file_controls_to_agent_state.py diff --git a/alembic/versions/c4eb5a907b38_add_file_controls_to_agent_state.py b/alembic/versions/c4eb5a907b38_add_file_controls_to_agent_state.py new file mode 100644 index 00000000..b9fa8426 --- /dev/null +++ b/alembic/versions/c4eb5a907b38_add_file_controls_to_agent_state.py @@ -0,0 +1,33 @@ +"""Add file controls to agent state + +Revision ID: c4eb5a907b38 +Revises: cce9a6174366 +Create Date: 2025-07-21 15:56:57.413000 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "c4eb5a907b38" +down_revision: Union[str, None] = "cce9a6174366" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("agents", sa.Column("max_files_open", sa.Integer(), nullable=True)) + op.add_column("agents", sa.Column("per_file_view_window_char_limit", sa.Integer(), nullable=True)) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("agents", "per_file_view_window_char_limit") + op.drop_column("agents", "max_files_open") + # ### end Alembic commands ### diff --git a/letta/constants.py b/letta/constants.py index 33c9745f..00b62388 100644 --- a/letta/constants.py +++ b/letta/constants.py @@ -326,7 +326,7 @@ MAX_ERROR_MESSAGE_CHAR_LIMIT = 500 CORE_MEMORY_PERSONA_CHAR_LIMIT: int = 5000 CORE_MEMORY_HUMAN_CHAR_LIMIT: int = 5000 CORE_MEMORY_BLOCK_CHAR_LIMIT: int = 5000 -CORE_MEMORY_SOURCE_CHAR_LIMIT: int = 50000 + # Function return limits FUNCTION_RETURN_CHAR_LIMIT = 6000 # ~300 words BASE_FUNCTION_RETURN_CHAR_LIMIT = 1000000 # very high (we rely on implementation) @@ -361,7 +361,9 @@ REDIS_DEFAULT_CACHE_PREFIX = "letta_cache" REDIS_RUN_ID_PREFIX = "agent:send_message:run_id" # TODO: This is temporary, eventually use token-based eviction -MAX_FILES_OPEN = 5 +# File based controls +DEFAULT_MAX_FILES_OPEN = 5 +DEFAULT_CORE_MEMORY_SOURCE_CHAR_LIMIT: int = 50000 GET_PROVIDERS_TIMEOUT_SECONDS = 10 diff --git a/letta/orm/agent.py b/letta/orm/agent.py index 44e8a6f0..bc2879d8 100644 --- a/letta/orm/agent.py +++ b/letta/orm/agent.py @@ -20,6 +20,7 @@ from letta.schemas.llm_config import LLMConfig from letta.schemas.memory import Memory from letta.schemas.response_format import ResponseFormatUnion from letta.schemas.tool_rule import ToolRule +from letta.utils import calculate_file_defaults_based_on_context_window if TYPE_CHECKING: from letta.orm.agents_tags import AgentsTags @@ -92,6 +93,14 @@ class Agent(SqlalchemyBase, OrganizationMixin, AsyncAttrs): # timezone timezone: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The timezone of the agent (for the context window).") + # file related controls + max_files_open: Mapped[Optional[int]] = mapped_column( + Integer, nullable=True, doc="Maximum number of files that can be open at once for this agent." + ) + per_file_view_window_char_limit: Mapped[Optional[int]] = mapped_column( + Integer, nullable=True, doc="The per-file view window character limit for this agent." + ) + # relationships organization: Mapped["Organization"] = relationship("Organization", back_populates="agents", lazy="raise") tool_exec_environment_variables: Mapped[List["AgentEnvironmentVariable"]] = relationship( @@ -146,6 +155,15 @@ class Agent(SqlalchemyBase, OrganizationMixin, AsyncAttrs): lazy="selectin", ) + def _get_per_file_view_window_char_limit(self) -> int: + """Get the per_file_view_window_char_limit, calculating defaults if None.""" + if self.per_file_view_window_char_limit is not None: + return self.per_file_view_window_char_limit + + context_window = self.llm_config.context_window if self.llm_config and self.llm_config.context_window else None + _, default_char_limit = calculate_file_defaults_based_on_context_window(context_window) + return default_char_limit + def to_pydantic(self, include_relationships: Optional[Set[str]] = None) -> PydanticAgentState: """ Converts the SQLAlchemy Agent model into its Pydantic counterpart. @@ -191,6 +209,8 @@ class Agent(SqlalchemyBase, OrganizationMixin, AsyncAttrs): "last_run_completion": self.last_run_completion, "last_run_duration_ms": self.last_run_duration_ms, "timezone": self.timezone, + "max_files_open": self.max_files_open, + "per_file_view_window_char_limit": self.per_file_view_window_char_limit, # optional field defaults "tags": [], "tools": [], @@ -208,7 +228,12 @@ class Agent(SqlalchemyBase, OrganizationMixin, AsyncAttrs): "sources": lambda: [s.to_pydantic() for s in self.sources], "memory": lambda: Memory( blocks=[b.to_pydantic() for b in self.core_memory], - file_blocks=[block for b in self.file_agents if (block := b.to_pydantic_block()) is not None], + file_blocks=[ + block + for b in self.file_agents + if (block := b.to_pydantic_block(per_file_view_window_char_limit=self._get_per_file_view_window_char_limit())) + is not None + ], prompt_template=get_prompt_template_for_agent_type(self.agent_type), ), "identity_ids": lambda: [i.id for i in self.identities], @@ -271,6 +296,8 @@ class Agent(SqlalchemyBase, OrganizationMixin, AsyncAttrs): "response_format": self.response_format, "last_run_completion": self.last_run_completion, "last_run_duration_ms": self.last_run_duration_ms, + "max_files_open": self.max_files_open, + "per_file_view_window_char_limit": self.per_file_view_window_char_limit, } optional_fields = { "tags": [], @@ -314,7 +341,11 @@ class Agent(SqlalchemyBase, OrganizationMixin, AsyncAttrs): state["sources"] = [s.to_pydantic() for s in sources] state["memory"] = Memory( blocks=[m.to_pydantic() for m in memory], - file_blocks=[block for b in file_agents if (block := b.to_pydantic_block()) is not None], + file_blocks=[ + block + for b in file_agents + if (block := b.to_pydantic_block(per_file_view_window_char_limit=self._get_per_file_view_window_char_limit())) is not None + ], prompt_template=get_prompt_template_for_agent_type(self.agent_type), ) state["identity_ids"] = [i.id for i in identities] diff --git a/letta/orm/files_agents.py b/letta/orm/files_agents.py index 208501e3..4ae6dd3b 100644 --- a/letta/orm/files_agents.py +++ b/letta/orm/files_agents.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Optional from sqlalchemy import Boolean, DateTime, ForeignKey, Index, String, Text, UniqueConstraint, func from sqlalchemy.orm import Mapped, mapped_column, relationship -from letta.constants import CORE_MEMORY_SOURCE_CHAR_LIMIT, FILE_IS_TRUNCATED_WARNING +from letta.constants import FILE_IS_TRUNCATED_WARNING from letta.orm.mixins import OrganizationMixin from letta.orm.sqlalchemy_base import SqlalchemyBase from letta.schemas.block import Block as PydanticBlock @@ -86,13 +86,13 @@ class FileAgent(SqlalchemyBase, OrganizationMixin): ) # TODO: This is temporary as we figure out if we want FileBlock as a first class citizen - def to_pydantic_block(self) -> PydanticBlock: + def to_pydantic_block(self, per_file_view_window_char_limit: int) -> PydanticBlock: visible_content = self.visible_content if self.visible_content and self.is_open else "" # Truncate content and add warnings here when converting from FileAgent to Block - if len(visible_content) > CORE_MEMORY_SOURCE_CHAR_LIMIT: + if len(visible_content) > per_file_view_window_char_limit: truncated_warning = f"...[TRUNCATED]\n{FILE_IS_TRUNCATED_WARNING}" - visible_content = visible_content[: CORE_MEMORY_SOURCE_CHAR_LIMIT - len(truncated_warning)] + visible_content = visible_content[: per_file_view_window_char_limit - len(truncated_warning)] visible_content += truncated_warning return PydanticBlock( @@ -100,5 +100,5 @@ class FileAgent(SqlalchemyBase, OrganizationMixin): label=self.file_name, # use denormalized file_name instead of self.file.file_name read_only=True, metadata={"source_id": self.source_id}, # use denormalized source_id - limit=CORE_MEMORY_SOURCE_CHAR_LIMIT, + limit=per_file_view_window_char_limit, ) diff --git a/letta/schemas/agent.py b/letta/schemas/agent.py index 2ffd4139..3da7bde2 100644 --- a/letta/schemas/agent.py +++ b/letta/schemas/agent.py @@ -19,7 +19,7 @@ from letta.schemas.response_format import ResponseFormatUnion from letta.schemas.source import Source from letta.schemas.tool import Tool from letta.schemas.tool_rule import ToolRule -from letta.utils import create_random_username +from letta.utils import calculate_file_defaults_based_on_context_window, create_random_username class AgentType(str, Enum): @@ -112,6 +112,16 @@ class AgentState(OrmMetadataBase, validate_assignment=True): # timezone timezone: Optional[str] = Field(None, description="The timezone of the agent (IANA format).") + # file related controls + max_files_open: Optional[int] = Field( + None, + description="Maximum number of files that can be open at once for this agent. Setting this too high may exceed the context window, which will break the agent.", + ) + per_file_view_window_char_limit: Optional[int] = Field( + None, + description="The per-file view window character limit for this agent. Setting this too high may exceed the context window, which will break the agent.", + ) + def get_agent_env_vars_as_dict(self) -> Dict[str, str]: # Get environment variables for this agent specifically per_agent_env_vars = {} @@ -119,6 +129,27 @@ class AgentState(OrmMetadataBase, validate_assignment=True): per_agent_env_vars[agent_env_var_obj.key] = agent_env_var_obj.value return per_agent_env_vars + @model_validator(mode="after") + def set_file_defaults_based_on_context_window(self) -> "AgentState": + """Set reasonable defaults for file-related fields based on the model's context window size.""" + # Only set defaults if not explicitly provided + if self.max_files_open is not None and self.per_file_view_window_char_limit is not None: + return self + + # Get context window size from llm_config + context_window = self.llm_config.context_window if self.llm_config and self.llm_config.context_window else None + + # Calculate defaults using the helper function + default_max_files, default_char_limit = calculate_file_defaults_based_on_context_window(context_window) + + # Apply defaults only if not set + if self.max_files_open is None: + self.max_files_open = default_max_files + if self.per_file_view_window_char_limit is None: + self.per_file_view_window_char_limit = default_char_limit + + return self + class CreateAgent(BaseModel, validate_assignment=True): # # all optional as server can generate defaults @@ -197,6 +228,14 @@ class CreateAgent(BaseModel, validate_assignment=True): # enable_sleeptime: Optional[bool] = Field(None, description="If set to True, memory management will move to a background agent thread.") response_format: Optional[ResponseFormatUnion] = Field(None, description="The response format for the agent.") timezone: Optional[str] = Field(None, description="The timezone of the agent (IANA format).") + max_files_open: Optional[int] = Field( + None, + description="Maximum number of files that can be open at once for this agent. Setting this too high may exceed the context window, which will break the agent.", + ) + per_file_view_window_char_limit: Optional[int] = Field( + None, + description="The per-file view window character limit for this agent. Setting this too high may exceed the context window, which will break the agent.", + ) @field_validator("name") @classmethod @@ -291,6 +330,14 @@ class UpdateAgent(BaseModel): last_run_completion: Optional[datetime] = Field(None, description="The timestamp when the agent last completed a run.") last_run_duration_ms: Optional[int] = Field(None, description="The duration in milliseconds of the agent's last run.") timezone: Optional[str] = Field(None, description="The timezone of the agent (IANA format).") + max_files_open: Optional[int] = Field( + None, + description="Maximum number of files that can be open at once for this agent. Setting this too high may exceed the context window, which will break the agent.", + ) + per_file_view_window_char_limit: Optional[int] = Field( + None, + description="The per-file view window character limit for this agent. Setting this too high may exceed the context window, which will break the agent.", + ) class Config: extra = "ignore" # Ignores extra fields diff --git a/letta/schemas/agent_file.py b/letta/schemas/agent_file.py index c9a18072..26b45cb0 100644 --- a/letta/schemas/agent_file.py +++ b/letta/schemas/agent_file.py @@ -145,6 +145,8 @@ class AgentSchema(CreateAgent): enable_sleeptime=False, # TODO: Need to figure out how to patch this response_format=agent_state.response_format, timezone=agent_state.timezone or "UTC", + max_files_open=agent_state.max_files_open, + per_file_view_window_char_limit=agent_state.per_file_view_window_char_limit, ) messages = await message_manager.list_messages_for_agent_async( diff --git a/letta/server/server.py b/letta/server/server.py index fcea1e8a..d86654b7 100644 --- a/letta/server/server.py +++ b/letta/server/server.py @@ -1427,6 +1427,7 @@ class SyncServer(Server): files_metadata=[file_metadata_with_content], visible_content_map=visible_content_map, actor=actor, + max_files_open=agent_state.max_files_open, ) for agent_state in agent_states ) @@ -1460,6 +1461,7 @@ class SyncServer(Server): files_metadata=file_metadata_with_content, visible_content_map=visible_content_map, actor=actor, + max_files_open=agent_state.max_files_open, ) if closed_files: diff --git a/letta/services/agent_file_manager.py b/letta/services/agent_file_manager.py index 03054656..e1691b4c 100644 --- a/letta/services/agent_file_manager.py +++ b/letta/services/agent_file_manager.py @@ -158,7 +158,11 @@ class AgentFileManager: for agent_state in agent_states: files_agents = await self.file_agent_manager.list_files_for_agent( - agent_id=agent_state.id, actor=actor, is_open_only=False, return_as_blocks=False + agent_id=agent_state.id, + actor=actor, + is_open_only=False, + return_as_blocks=False, + per_file_view_window_char_limit=agent_state.per_file_view_window_char_limit, ) # cache the results for reuse during conversion if files_agents_cache is not None: @@ -182,7 +186,11 @@ class AgentFileManager: files_agents = files_agents_cache[agent_state.id] else: files_agents = await self.file_agent_manager.list_files_for_agent( - agent_id=agent_state.id, actor=actor, is_open_only=False, return_as_blocks=False + agent_id=agent_state.id, + actor=actor, + is_open_only=False, + return_as_blocks=False, + per_file_view_window_char_limit=agent_state.per_file_view_window_char_limit, ) agent_schema = await AgentSchema.from_agent_state( agent_state, message_manager=self.message_manager, files_agents=files_agents, actor=actor @@ -510,7 +518,11 @@ class AgentFileManager: # Bulk attach files to agent await self.file_agent_manager.attach_files_bulk( - agent_id=agent_db_id, files_metadata=files_for_agent, visible_content_map=visible_content_map, actor=actor + agent_id=agent_db_id, + files_metadata=files_for_agent, + visible_content_map=visible_content_map, + actor=actor, + max_files_open=agent_schema.max_files_open, ) imported_count += len(files_for_agent) diff --git a/letta/services/agent_manager.py b/letta/services/agent_manager.py index 8d3e6457..920d1810 100644 --- a/letta/services/agent_manager.py +++ b/letta/services/agent_manager.py @@ -359,6 +359,8 @@ class AgentManager: created_by_id=actor.id, last_updated_by_id=actor.id, timezone=agent_create.timezone, + max_files_open=agent_create.max_files_open, + per_file_view_window_char_limit=agent_create.per_file_view_window_char_limit, ) if _test_only_force_id: @@ -548,6 +550,8 @@ class AgentManager: created_by_id=actor.id, last_updated_by_id=actor.id, timezone=agent_create.timezone if agent_create.timezone else DEFAULT_TIMEZONE, + max_files_open=agent_create.max_files_open, + per_file_view_window_char_limit=agent_create.per_file_view_window_char_limit, ) if _test_only_force_id: @@ -711,6 +715,9 @@ class AgentManager: "response_format": agent_update.response_format, "last_run_completion": agent_update.last_run_completion, "last_run_duration_ms": agent_update.last_run_duration_ms, + "max_files_open": agent_update.max_files_open, + "per_file_view_window_char_limit": agent_update.per_file_view_window_char_limit, + "timezone": agent_update.timezone, } for col, val in scalar_updates.items(): if val is not None: @@ -834,6 +841,8 @@ class AgentManager: "last_run_completion": agent_update.last_run_completion, "last_run_duration_ms": agent_update.last_run_duration_ms, "timezone": agent_update.timezone, + "max_files_open": agent_update.max_files_open, + "per_file_view_window_char_limit": agent_update.per_file_view_window_char_limit, } for col, val in scalar_updates.items(): if val is not None: @@ -1864,7 +1873,10 @@ class AgentManager: if file_block_names: file_blocks = await self.file_agent_manager.get_all_file_blocks_by_name( - file_names=file_block_names, agent_id=agent_state.id, actor=actor + file_names=file_block_names, + agent_id=agent_state.id, + actor=actor, + per_file_view_window_char_limit=agent_state.per_file_view_window_char_limit, ) agent_state.memory.file_blocks = [b for b in file_blocks if b is not None] @@ -1873,7 +1885,12 @@ class AgentManager: @enforce_types @trace_method async def refresh_file_blocks(self, agent_state: PydanticAgentState, actor: PydanticUser) -> PydanticAgentState: - file_blocks = await self.file_agent_manager.list_files_for_agent(agent_id=agent_state.id, actor=actor, return_as_blocks=True) + file_blocks = await self.file_agent_manager.list_files_for_agent( + agent_id=agent_state.id, + per_file_view_window_char_limit=agent_state.per_file_view_window_char_limit, + actor=actor, + return_as_blocks=True, + ) agent_state.memory.file_blocks = [b for b in file_blocks if b is not None] return agent_state @@ -2862,6 +2879,64 @@ class AgentManager: results = [row[0] for row in result.all()] return results + @enforce_types + @trace_method + async def get_agent_max_files_open_async(self, agent_id: str, actor: PydanticUser) -> int: + """Get max_files_open for an agent. + + This is a performant query that only fetches the specific field needed. + + Args: + agent_id: The ID of the agent + actor: The user making the request + + Returns: + max_files_open value + """ + async with db_registry.async_session() as session: + # Direct query for just max_files_open + result = await session.execute( + select(AgentModel.max_files_open) + .where(AgentModel.id == agent_id) + .where(AgentModel.organization_id == actor.organization_id) + .where(AgentModel.is_deleted == False) + ) + row = result.scalar_one_or_none() + + if row is None: + raise ValueError(f"Agent {agent_id} not found") + + return row + + @enforce_types + @trace_method + async def get_agent_per_file_view_window_char_limit_async(self, agent_id: str, actor: PydanticUser) -> int: + """Get per_file_view_window_char_limit for an agent. + + This is a performant query that only fetches the specific field needed. + + Args: + agent_id: The ID of the agent + actor: The user making the request + + Returns: + per_file_view_window_char_limit value + """ + async with db_registry.async_session() as session: + # Direct query for just per_file_view_window_char_limit + result = await session.execute( + select(AgentModel.per_file_view_window_char_limit) + .where(AgentModel.id == agent_id) + .where(AgentModel.organization_id == actor.organization_id) + .where(AgentModel.is_deleted == False) + ) + row = result.scalar_one_or_none() + + if row is None: + raise ValueError(f"Agent {agent_id} not found") + + return row + @trace_method async def get_context_window(self, agent_id: str, actor: PydanticUser) -> ContextWindowOverview: agent_state, system_message, num_messages, num_archival_memories = await self.rebuild_system_prompt_async( diff --git a/letta/services/files_agents_manager.py b/letta/services/files_agents_manager.py index 9d8cc838..38b98138 100644 --- a/letta/services/files_agents_manager.py +++ b/letta/services/files_agents_manager.py @@ -3,7 +3,6 @@ from typing import List, Optional from sqlalchemy import and_, func, select, update -from letta.constants import MAX_FILES_OPEN from letta.log import get_logger from letta.orm.errors import NoResultFound from letta.orm.files_agents import FileAgent as FileAgentModel @@ -31,6 +30,7 @@ class FileAgentManager: file_name: str, source_id: str, actor: PydanticUser, + max_files_open: int, is_open: bool = True, visible_content: Optional[str] = None, ) -> tuple[PydanticFileAgent, List[str]]: @@ -40,7 +40,7 @@ class FileAgentManager: • If the row already exists → update `is_open`, `visible_content` and always refresh `last_accessed_at`. • Otherwise create a brand-new association. - • If is_open=True, enforces MAX_FILES_OPEN using LRU eviction. + • If is_open=True, enforces max_files_open using LRU eviction. Returns: Tuple of (file_agent, closed_file_names) @@ -54,6 +54,7 @@ class FileAgentManager: source_id=source_id, actor=actor, visible_content=visible_content or "", + max_files_open=max_files_open, ) # Get the updated file agent to return @@ -177,6 +178,7 @@ class FileAgentManager: *, file_names: List[str], agent_id: str, + per_file_view_window_char_limit: int, actor: PydanticUser, ) -> List[PydanticBlock]: """ @@ -185,6 +187,7 @@ class FileAgentManager: Args: file_names: List of file names to retrieve agent_id: ID of the agent to retrieve file blocks for + per_file_view_window_char_limit: The per-file view window char limit actor: The user making the request Returns: @@ -207,7 +210,7 @@ class FileAgentManager: rows = (await session.execute(query)).scalars().all() # Convert to Pydantic models - return [row.to_pydantic_block() for row in rows] + return [row.to_pydantic_block(per_file_view_window_char_limit=per_file_view_window_char_limit) for row in rows] @enforce_types @trace_method @@ -222,7 +225,12 @@ class FileAgentManager: @enforce_types @trace_method async def list_files_for_agent( - self, agent_id: str, actor: PydanticUser, is_open_only: bool = False, return_as_blocks: bool = False + self, + agent_id: str, + per_file_view_window_char_limit: int, + actor: PydanticUser, + is_open_only: bool = False, + return_as_blocks: bool = False, ) -> List[PydanticFileAgent]: """Return associations for *agent_id* (filtering by `is_open` if asked).""" async with db_registry.async_session() as session: @@ -236,7 +244,7 @@ class FileAgentManager: rows = (await session.execute(select(FileAgentModel).where(and_(*conditions)))).scalars().all() if return_as_blocks: - return [r.to_pydantic_block() for r in rows] + return [r.to_pydantic_block(per_file_view_window_char_limit=per_file_view_window_char_limit) for r in rows] else: return [r.to_pydantic() for r in rows] @@ -334,7 +342,7 @@ class FileAgentManager: @enforce_types @trace_method async def enforce_max_open_files_and_open( - self, *, agent_id: str, file_id: str, file_name: str, source_id: str, actor: PydanticUser, visible_content: str + self, *, agent_id: str, file_id: str, file_name: str, source_id: str, actor: PydanticUser, visible_content: str, max_files_open: int ) -> tuple[List[str], bool]: """ Efficiently handle LRU eviction and file opening in a single transaction. @@ -386,7 +394,7 @@ class FileAgentManager: # Calculate how many files need to be closed current_other_count = len(other_open_files) - target_other_count = MAX_FILES_OPEN - 1 # Reserve 1 slot for file we're opening + target_other_count = max_files_open - 1 # Reserve 1 slot for file we're opening closed_file_names = [] if current_other_count > target_other_count: @@ -443,6 +451,7 @@ class FileAgentManager: *, agent_id: str, files_metadata: list[FileMetadata], + max_files_open: int, visible_content_map: Optional[dict[str, str]] = None, actor: PydanticUser, ) -> list[str]: @@ -496,17 +505,17 @@ class FileAgentManager: still_open_names = [r.file_name for r in currently_open if r.file_name not in new_names_set] # decide final open set - if len(new_names) >= MAX_FILES_OPEN: - final_open = new_names[:MAX_FILES_OPEN] + if len(new_names) >= max_files_open: + final_open = new_names[:max_files_open] else: - room_for_old = MAX_FILES_OPEN - len(new_names) + room_for_old = max_files_open - len(new_names) final_open = new_names + still_open_names[-room_for_old:] final_open_set = set(final_open) closed_file_names = [r.file_name for r in currently_open if r.file_name not in final_open_set] - # Add new files that won't be opened due to MAX_FILES_OPEN limit - if len(new_names) >= MAX_FILES_OPEN: - closed_file_names.extend(new_names[MAX_FILES_OPEN:]) + # Add new files that won't be opened due to max_files_open limit + if len(new_names) >= max_files_open: + closed_file_names.extend(new_names[max_files_open:]) evicted_ids = [r.file_id for r in currently_open if r.file_name in closed_file_names] # upsert requested files diff --git a/letta/services/helpers/agent_manager_helper.py b/letta/services/helpers/agent_manager_helper.py index eff7ff58..e5ff86a4 100644 --- a/letta/services/helpers/agent_manager_helper.py +++ b/letta/services/helpers/agent_manager_helper.py @@ -22,11 +22,12 @@ from letta.constants import ( from letta.embeddings import embedding_model from letta.helpers import ToolRulesSolver from letta.helpers.datetime_helpers import format_datetime, get_local_time, get_local_time_fast -from letta.orm import AgentPassage, SourcePassage, SourcesAgents from letta.orm.agent import Agent as AgentModel from letta.orm.agents_tags import AgentsTags from letta.orm.errors import NoResultFound from letta.orm.identity import Identity +from letta.orm.passage import AgentPassage, SourcePassage +from letta.orm.sources_agents import SourcesAgents from letta.orm.sqlite_functions import adapt_array from letta.otel.tracing import trace_method from letta.prompts import gpt_system @@ -45,7 +46,7 @@ from letta.system import get_initial_boot_messages, get_login_event, package_fun # Static methods @trace_method def _process_relationship( - session, agent: AgentModel, relationship_name: str, model_class, item_ids: List[str], allow_partial=False, replace=True + session, agent: "AgentModel", relationship_name: str, model_class, item_ids: List[str], allow_partial=False, replace=True ): """ Generalized function to handle relationships like tools, sources, and blocks using item IDs. @@ -88,7 +89,7 @@ def _process_relationship( @trace_method async def _process_relationship_async( - session, agent: AgentModel, relationship_name: str, model_class, item_ids: List[str], allow_partial=False, replace=True + session, agent: "AgentModel", relationship_name: str, model_class, item_ids: List[str], allow_partial=False, replace=True ): """ Generalized function to handle relationships like tools, sources, and blocks using item IDs. @@ -130,7 +131,7 @@ async def _process_relationship_async( current_relationship.extend(new_items) -def _process_tags(agent: AgentModel, tags: List[str], replace=True): +def _process_tags(agent: "AgentModel", tags: List[str], replace=True): """ Handles tags for an agent. diff --git a/letta/services/tool_executor/files_tool_executor.py b/letta/services/tool_executor/files_tool_executor.py index b56b2253..6e5a6304 100644 --- a/letta/services/tool_executor/files_tool_executor.py +++ b/letta/services/tool_executor/files_tool_executor.py @@ -2,7 +2,7 @@ import asyncio import re from typing import Any, Dict, List, Optional -from letta.constants import MAX_FILES_OPEN, PINECONE_TEXT_FIELD_NAME +from letta.constants import PINECONE_TEXT_FIELD_NAME from letta.functions.types import FileOpenRequest from letta.helpers.pinecone_utils import search_pinecone_index, should_use_pinecone from letta.log import get_logger @@ -117,8 +117,10 @@ class LettaFileToolExecutor(ToolExecutor): file_requests = parsed_requests # Validate file count first - if len(file_requests) > MAX_FILES_OPEN: - raise ValueError(f"Cannot open {len(file_requests)} files: exceeds maximum limit of {MAX_FILES_OPEN} files") + if len(file_requests) > agent_state.max_files_open: + raise ValueError( + f"Cannot open {len(file_requests)} files: exceeds configured maximum limit of {agent_state.max_files_open} files" + ) if not file_requests: raise ValueError("No file requests provided") @@ -186,6 +188,7 @@ class LettaFileToolExecutor(ToolExecutor): source_id=file.source_id, actor=self.actor, visible_content=visible_content, + max_files_open=agent_state.max_files_open, ) opened_files.append(file_name) @@ -329,7 +332,9 @@ class LettaFileToolExecutor(ToolExecutor): include_regex = re.compile(include_pattern, re.IGNORECASE) # Get all attached files for this agent - file_agents = await self.files_agents_manager.list_files_for_agent(agent_id=agent_state.id, actor=self.actor) + file_agents = await self.files_agents_manager.list_files_for_agent( + agent_id=agent_state.id, per_file_view_window_char_limit=agent_state.per_file_view_window_char_limit, actor=self.actor + ) if not file_agents: return "No files are currently attached to search" @@ -509,7 +514,9 @@ class LettaFileToolExecutor(ToolExecutor): return f"No valid source IDs found for attached files" # Get all attached files for this agent - file_agents = await self.files_agents_manager.list_files_for_agent(agent_id=agent_state.id, actor=self.actor) + file_agents = await self.files_agents_manager.list_files_for_agent( + agent_id=agent_state.id, per_file_view_window_char_limit=agent_state.per_file_view_window_char_limit, actor=self.actor + ) if not file_agents: return "No files are currently attached to search" diff --git a/letta/utils.py b/letta/utils.py index 29699bd6..2b8a4501 100644 --- a/letta/utils.py +++ b/letta/utils.py @@ -17,7 +17,7 @@ from contextlib import contextmanager from datetime import datetime, timezone from functools import wraps from logging import Logger -from typing import Any, Coroutine, Union, _GenericAlias, get_args, get_origin, get_type_hints +from typing import Any, Coroutine, Optional, Union, _GenericAlias, get_args, get_origin, get_type_hints from urllib.parse import urljoin, urlparse import demjson3 as demjson @@ -30,6 +30,8 @@ from letta.constants import ( CLI_WARNING_PREFIX, CORE_MEMORY_HUMAN_CHAR_LIMIT, CORE_MEMORY_PERSONA_CHAR_LIMIT, + DEFAULT_CORE_MEMORY_SOURCE_CHAR_LIMIT, + DEFAULT_MAX_FILES_OPEN, ERROR_MESSAGE_PREFIX, LETTA_DIR, MAX_FILENAME_LENGTH, @@ -1206,3 +1208,34 @@ async def get_latest_alembic_revision() -> str: except Exception as e: logger.error(f"Error getting latest alembic revision: {e}") return "unknown" + + +def calculate_file_defaults_based_on_context_window(context_window: Optional[int]) -> tuple[int, int]: + """Calculate reasonable defaults for max_files_open and per_file_view_window_char_limit + based on the model's context window size. + + Args: + context_window: The context window size of the model. If None, returns conservative defaults. + + Returns: + A tuple of (max_files_open, per_file_view_window_char_limit) + """ + if not context_window: + # If no context window info, use conservative defaults + return DEFAULT_MAX_FILES_OPEN, DEFAULT_CORE_MEMORY_SOURCE_CHAR_LIMIT + + # Define defaults based on context window ranges + # Assuming ~4 chars per token + # Available chars = available_tokens * 4 + + # TODO: Check my math here + if context_window <= 8_000: # Small models (4K-8K) + return 3, 5_000 # ~3.75K tokens + elif context_window <= 32_000: # Medium models (16K-32K) + return 5, 15_000 # ~18.75K tokens + elif context_window <= 128_000: # Large models (100K-128K) + return 10, 25_000 # ~62.5K tokens + elif context_window <= 200_000: # Very large models (128K-200K) + return 10, 50_000 # ~128k tokens + else: # Extremely large models (200K+) + return 15, 50_000 # ~187.5k tokens diff --git a/tests/test_managers.py b/tests/test_managers.py index ffc583c0..7eb5307d 100644 --- a/tests/test_managers.py +++ b/tests/test_managers.py @@ -88,6 +88,7 @@ from letta.services.block_manager import BlockManager from letta.services.helpers.agent_manager_helper import calculate_base_tools, calculate_multi_agent_tools, validate_agent_exists_async from letta.services.step_manager import FeedbackType from letta.settings import tool_settings +from letta.utils import calculate_file_defaults_based_on_context_window from tests.helpers.utils import comprehensive_agent_checks, validate_context_window_overview from tests.utils import random_string @@ -686,6 +687,7 @@ async def file_attachment(server, default_user, sarah_agent, default_file): source_id=default_file.source_id, actor=default_user, visible_content="initial", + max_files_open=sarah_agent.max_files_open, ) yield assoc @@ -933,6 +935,7 @@ async def test_get_context_window_basic( source_id=default_file.source_id, actor=default_user, visible_content="hello", + max_files_open=created_agent.max_files_open, ) # Get context window and check for basic appearances @@ -1057,6 +1060,113 @@ async def test_update_agent( assert updated_agent.updated_at > last_updated_timestamp +@pytest.mark.asyncio +async def test_agent_file_defaults_based_on_context_window(server: SyncServer, default_user, default_block, event_loop): + """Test that file-related defaults are set based on the model's context window size""" + + # test with small context window model (8k) + llm_config_small = LLMConfig.default_config("gpt-4o-mini") + llm_config_small.context_window = 8000 + create_agent_request = CreateAgent( + name="test_agent_small_context", + llm_config=llm_config_small, + embedding_config=EmbeddingConfig.default_config(provider="openai"), + block_ids=[default_block.id], + include_base_tools=False, + ) + agent_state = await server.agent_manager.create_agent_async( + create_agent_request, + actor=default_user, + ) + assert agent_state.max_files_open == 3 + assert ( + agent_state.per_file_view_window_char_limit == calculate_file_defaults_based_on_context_window(llm_config_small.context_window)[1] + ) + server.agent_manager.delete_agent(agent_id=agent_state.id, actor=default_user) + + # test with medium context window model (32k) + llm_config_medium = LLMConfig.default_config("gpt-4o-mini") + llm_config_medium.context_window = 32000 + create_agent_request = CreateAgent( + name="test_agent_medium_context", + llm_config=llm_config_medium, + embedding_config=EmbeddingConfig.default_config(provider="openai"), + block_ids=[default_block.id], + include_base_tools=False, + ) + agent_state = await server.agent_manager.create_agent_async( + create_agent_request, + actor=default_user, + ) + assert agent_state.max_files_open == 5 + assert ( + agent_state.per_file_view_window_char_limit == calculate_file_defaults_based_on_context_window(llm_config_medium.context_window)[1] + ) + server.agent_manager.delete_agent(agent_id=agent_state.id, actor=default_user) + + # test with large context window model (128k) + llm_config_large = LLMConfig.default_config("gpt-4o-mini") + llm_config_large.context_window = 128000 + create_agent_request = CreateAgent( + name="test_agent_large_context", + llm_config=llm_config_large, + embedding_config=EmbeddingConfig.default_config(provider="openai"), + block_ids=[default_block.id], + include_base_tools=False, + ) + agent_state = await server.agent_manager.create_agent_async( + create_agent_request, + actor=default_user, + ) + assert agent_state.max_files_open == 10 + assert ( + agent_state.per_file_view_window_char_limit == calculate_file_defaults_based_on_context_window(llm_config_large.context_window)[1] + ) + server.agent_manager.delete_agent(agent_id=agent_state.id, actor=default_user) + + +@pytest.mark.asyncio +async def test_agent_file_defaults_explicit_values(server: SyncServer, default_user, default_block, event_loop): + """Test that explicitly set file-related values are respected""" + + llm_config_explicit = LLMConfig.default_config("gpt-4o-mini") + llm_config_explicit.context_window = 32000 # would normally get defaults of 5 and 30k + create_agent_request = CreateAgent( + name="test_agent_explicit_values", + llm_config=llm_config_explicit, + embedding_config=EmbeddingConfig.default_config(provider="openai"), + block_ids=[default_block.id], + include_base_tools=False, + max_files_open=20, # explicit value + per_file_view_window_char_limit=500_000, # explicit value + ) + agent_state = await server.agent_manager.create_agent_async( + create_agent_request, + actor=default_user, + ) + # verify explicit values are used instead of defaults + assert agent_state.max_files_open == 20 + assert agent_state.per_file_view_window_char_limit == 500_000 + server.agent_manager.delete_agent(agent_id=agent_state.id, actor=default_user) + + +@pytest.mark.asyncio +async def test_update_agent_file_fields(server: SyncServer, comprehensive_test_agent_fixture, default_user, event_loop): + """Test updating file-related fields on an existing agent""" + + agent, _ = comprehensive_test_agent_fixture + + # update file-related fields + update_request = UpdateAgent( + max_files_open=15, + per_file_view_window_char_limit=150_000, + ) + updated_agent = await server.agent_manager.update_agent_async(agent.id, update_request, actor=default_user) + + assert updated_agent.max_files_open == 15 + assert updated_agent.per_file_view_window_char_limit == 150_000 + + # ====================================================================================================================== # AgentManager Tests - Listing # ====================================================================================================================== @@ -7654,6 +7764,7 @@ async def test_attach_creates_association(server, default_user, sarah_agent, def source_id=default_file.source_id, actor=default_user, visible_content="hello", + max_files_open=sarah_agent.max_files_open, ) assert assoc.agent_id == sarah_agent.id @@ -7677,6 +7788,7 @@ async def test_attach_is_idempotent(server, default_user, sarah_agent, default_f source_id=default_file.source_id, actor=default_user, visible_content="first", + max_files_open=sarah_agent.max_files_open, ) # second attach with different params @@ -7688,6 +7800,7 @@ async def test_attach_is_idempotent(server, default_user, sarah_agent, default_f actor=default_user, is_open=False, visible_content="second", + max_files_open=sarah_agent.max_files_open, ) assert a1.id == a2.id @@ -7764,6 +7877,7 @@ async def test_list_files_and_agents( file_name=default_file.file_name, source_id=default_file.source_id, actor=default_user, + max_files_open=charles_agent.max_files_open, ) # default_file ↔ sarah (open) await server.file_agent_manager.attach_file( @@ -7772,6 +7886,7 @@ async def test_list_files_and_agents( file_name=default_file.file_name, source_id=default_file.source_id, actor=default_user, + max_files_open=sarah_agent.max_files_open, ) # another_file ↔ sarah (closed) await server.file_agent_manager.attach_file( @@ -7781,12 +7896,17 @@ async def test_list_files_and_agents( source_id=another_file.source_id, actor=default_user, is_open=False, + max_files_open=sarah_agent.max_files_open, ) - files_for_sarah = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user) + files_for_sarah = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user + ) assert {f.file_id for f in files_for_sarah} == {default_file.id, another_file.id} - open_only = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user, is_open_only=True) + open_only = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True + ) assert {f.file_id for f in open_only} == {default_file.id} agents_for_default = await server.file_agent_manager.list_agents_for_file(default_file.id, actor=default_user) @@ -7832,10 +7952,13 @@ async def test_org_scoping( file_name=default_file.file_name, source_id=default_file.source_id, actor=default_user, + max_files_open=sarah_agent.max_files_open, ) # other org should see nothing - files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=other_user_different_org) + files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=other_user_different_org + ) assert files == [] @@ -7870,6 +7993,7 @@ async def test_mark_access_bulk(server, default_user, sarah_agent, default_sourc source_id=file.source_id, actor=default_user, visible_content=f"content for {file.file_name}", + max_files_open=sarah_agent.max_files_open, ) attached_files.append(file_agent) @@ -7898,14 +8022,15 @@ async def test_mark_access_bulk(server, default_user, sarah_agent, default_sourc @pytest.mark.asyncio async def test_lru_eviction_on_attach(server, default_user, sarah_agent, default_source): - """Test that attaching files beyond MAX_FILES_OPEN triggers LRU eviction.""" + """Test that attaching files beyond max_files_open triggers LRU eviction.""" import time - from letta.constants import MAX_FILES_OPEN + # Use the agent's configured max_files_open + max_files_open = sarah_agent.max_files_open # Create more files than the limit files = [] - for i in range(MAX_FILES_OPEN + 2): # 7 files for MAX_FILES_OPEN=5 + for i in range(max_files_open + 2): # e.g., 7 files for max_files_open=5 file_metadata = PydanticFileMetadata( file_name=f"lru_test_file_{i}.txt", organization_id=default_user.organization_id, @@ -7929,41 +8054,52 @@ async def test_lru_eviction_on_attach(server, default_user, sarah_agent, default source_id=file.source_id, actor=default_user, visible_content=f"content for {file.file_name}", + max_files_open=sarah_agent.max_files_open, ) attached_files.append(file_agent) all_closed_files.extend(closed_files) - # Check that we never exceed MAX_FILES_OPEN - open_files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user, is_open_only=True) - assert len(open_files) <= MAX_FILES_OPEN, f"Should never exceed {MAX_FILES_OPEN} open files" + # Check that we never exceed max_files_open + open_files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, + per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, + actor=default_user, + is_open_only=True, + ) + assert len(open_files) <= max_files_open, f"Should never exceed {max_files_open} open files" - # Should have closed exactly 2 files (7 - 5 = 2) - assert len(all_closed_files) == 2, f"Should have closed 2 files, but closed: {all_closed_files}" + # Should have closed exactly 2 files (e.g., 7 - 5 = 2 for max_files_open=5) + expected_closed_count = len(files) - max_files_open + assert ( + len(all_closed_files) == expected_closed_count + ), f"Should have closed {expected_closed_count} files, but closed: {all_closed_files}" - # Check that the oldest files were closed (first 2 files attached) - expected_closed = [files[0].file_name, files[1].file_name] + # Check that the oldest files were closed (first N files attached) + expected_closed = [files[i].file_name for i in range(expected_closed_count)] assert set(all_closed_files) == set(expected_closed), f"Wrong files closed. Expected {expected_closed}, got {all_closed_files}" - # Check that exactly MAX_FILES_OPEN files are open - open_files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user, is_open_only=True) - assert len(open_files) == MAX_FILES_OPEN + # Check that exactly max_files_open files are open + open_files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True + ) + assert len(open_files) == max_files_open # Check that the most recently attached files are still open open_file_names = {f.file_name for f in open_files} - expected_open = {files[i].file_name for i in range(2, MAX_FILES_OPEN + 2)} # files 2-6 + expected_open = {files[i].file_name for i in range(expected_closed_count, len(files))} # last max_files_open files assert open_file_names == expected_open @pytest.mark.asyncio async def test_lru_eviction_on_open_file(server, default_user, sarah_agent, default_source): - """Test that opening a file beyond MAX_FILES_OPEN triggers LRU eviction.""" + """Test that opening a file beyond max_files_open triggers LRU eviction.""" import time - from letta.constants import MAX_FILES_OPEN + max_files_open = sarah_agent.max_files_open # Create files equal to the limit files = [] - for i in range(MAX_FILES_OPEN + 1): # 6 files for MAX_FILES_OPEN=5 + for i in range(max_files_open + 1): # 6 files for max_files_open=5 file_metadata = PydanticFileMetadata( file_name=f"open_test_file_{i}.txt", organization_id=default_user.organization_id, @@ -7972,8 +8108,8 @@ async def test_lru_eviction_on_open_file(server, default_user, sarah_agent, defa file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text=f"test content {i}") files.append(file) - # Attach first MAX_FILES_OPEN files - for i in range(MAX_FILES_OPEN): + # Attach first max_files_open files + for i in range(max_files_open): time.sleep(0.1) # Small delay for different timestamps await server.file_agent_manager.attach_file( agent_id=sarah_agent.id, @@ -7982,6 +8118,7 @@ async def test_lru_eviction_on_open_file(server, default_user, sarah_agent, defa source_id=files[i].source_id, actor=default_user, visible_content=f"content for {files[i].file_name}", + max_files_open=sarah_agent.max_files_open, ) # Attach the last file as closed @@ -7993,13 +8130,18 @@ async def test_lru_eviction_on_open_file(server, default_user, sarah_agent, defa actor=default_user, is_open=False, visible_content=f"content for {files[-1].file_name}", + max_files_open=sarah_agent.max_files_open, ) - # All files should be attached but only MAX_FILES_OPEN should be open - all_files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user) - open_files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user, is_open_only=True) - assert len(all_files) == MAX_FILES_OPEN + 1 - assert len(open_files) == MAX_FILES_OPEN + # All files should be attached but only max_files_open should be open + all_files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user + ) + open_files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True + ) + assert len(all_files) == max_files_open + 1 + assert len(open_files) == max_files_open # Wait a moment time.sleep(0.1) @@ -8012,15 +8154,18 @@ async def test_lru_eviction_on_open_file(server, default_user, sarah_agent, defa source_id=files[-1].source_id, actor=default_user, visible_content="updated content", + max_files_open=sarah_agent.max_files_open, ) # Should have closed 1 file (the oldest one) assert len(closed_files) == 1, f"Should have closed 1 file, got: {closed_files}" assert closed_files[0] == files[0].file_name, f"Should have closed oldest file {files[0].file_name}" - # Check that exactly MAX_FILES_OPEN files are still open - open_files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user, is_open_only=True) - assert len(open_files) == MAX_FILES_OPEN + # Check that exactly max_files_open files are still open + open_files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True + ) + assert len(open_files) == max_files_open # Check that the newly opened file is open and the oldest is closed last_file_agent = await server.file_agent_manager.get_file_agent_by_id( @@ -8039,11 +8184,11 @@ async def test_lru_no_eviction_when_reopening_same_file(server, default_user, sa """Test that reopening an already open file doesn't trigger unnecessary eviction.""" import time - from letta.constants import MAX_FILES_OPEN + max_files_open = sarah_agent.max_files_open # Create files equal to the limit files = [] - for i in range(MAX_FILES_OPEN): + for i in range(max_files_open): file_metadata = PydanticFileMetadata( file_name=f"reopen_test_file_{i}.txt", organization_id=default_user.organization_id, @@ -8062,11 +8207,14 @@ async def test_lru_no_eviction_when_reopening_same_file(server, default_user, sa source_id=file.source_id, actor=default_user, visible_content=f"content for {file.file_name}", + max_files_open=sarah_agent.max_files_open, ) # All files should be open - open_files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user, is_open_only=True) - assert len(open_files) == MAX_FILES_OPEN + open_files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True + ) + assert len(open_files) == max_files_open initial_open_names = {f.file_name for f in open_files} # Wait a moment @@ -8080,6 +8228,7 @@ async def test_lru_no_eviction_when_reopening_same_file(server, default_user, sa source_id=files[-1].source_id, actor=default_user, visible_content="updated content", + max_files_open=sarah_agent.max_files_open, ) # Should not have closed any files since we're within the limit @@ -8087,8 +8236,10 @@ async def test_lru_no_eviction_when_reopening_same_file(server, default_user, sa assert was_already_open is True, "File should have been detected as already open" # All the same files should still be open - open_files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user, is_open_only=True) - assert len(open_files) == MAX_FILES_OPEN + open_files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True + ) + assert len(open_files) == max_files_open final_open_names = {f.file_name for f in open_files} assert initial_open_names == final_open_names, "Same files should remain open" @@ -8113,6 +8264,7 @@ async def test_last_accessed_at_updates_correctly(server, default_user, sarah_ag source_id=file.source_id, actor=default_user, visible_content="initial content", + max_files_open=sarah_agent.max_files_open, ) initial_time = file_agent.last_accessed_at @@ -8166,13 +8318,16 @@ async def test_attach_files_bulk_basic(server, default_user, sarah_agent, defaul files_metadata=files, visible_content_map=visible_content_map, actor=default_user, + max_files_open=sarah_agent.max_files_open, ) # Should not close any files since we're under the limit assert closed_files == [] # Verify all files are attached and open - attached_files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user, is_open_only=True) + attached_files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True + ) assert len(attached_files) == 3 attached_file_names = {f.file_name for f in attached_files} @@ -8213,10 +8368,13 @@ async def test_attach_files_bulk_deduplication(server, default_user, sarah_agent files_metadata=files_to_attach, visible_content_map=visible_content_map, actor=default_user, + max_files_open=sarah_agent.max_files_open, ) # Should only attach one file (deduplicated) - attached_files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user) + attached_files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user + ) assert len(attached_files) == 1 assert attached_files[0].file_name == "duplicate_test.txt" @@ -8226,11 +8384,11 @@ async def test_attach_files_bulk_lru_eviction(server, default_user, sarah_agent, """Test that attach_files_bulk properly handles LRU eviction without duplicates.""" import time - from letta.constants import MAX_FILES_OPEN + max_files_open = sarah_agent.max_files_open # First, fill up to the max with individual files existing_files = [] - for i in range(MAX_FILES_OPEN): + for i in range(max_files_open): file_metadata = PydanticFileMetadata( file_name=f"existing_{i}.txt", organization_id=default_user.organization_id, @@ -8247,11 +8405,14 @@ async def test_attach_files_bulk_lru_eviction(server, default_user, sarah_agent, source_id=file.source_id, actor=default_user, visible_content=f"existing content {i}", + max_files_open=sarah_agent.max_files_open, ) # Verify we're at the limit - open_files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user, is_open_only=True) - assert len(open_files) == MAX_FILES_OPEN + open_files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True + ) + assert len(open_files) == max_files_open # Now bulk attach 3 new files (should trigger LRU eviction) new_files = [] @@ -8272,6 +8433,7 @@ async def test_attach_files_bulk_lru_eviction(server, default_user, sarah_agent, files_metadata=new_files, visible_content_map=visible_content_map, actor=default_user, + max_files_open=sarah_agent.max_files_open, ) # Should have closed exactly 3 files (oldest ones) @@ -8285,9 +8447,11 @@ async def test_attach_files_bulk_lru_eviction(server, default_user, sarah_agent, actual_closed = set(closed_files) assert actual_closed == expected_closed - # Verify we still have exactly MAX_FILES_OPEN files open - open_files_after = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user, is_open_only=True) - assert len(open_files_after) == MAX_FILES_OPEN + # Verify we still have exactly max_files_open files open + open_files_after = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True + ) + assert len(open_files_after) == max_files_open # Verify the new files are open open_file_names = {f.file_name for f in open_files_after} @@ -8314,6 +8478,7 @@ async def test_attach_files_bulk_mixed_existing_new(server, default_user, sarah_ actor=default_user, visible_content="old content", is_open=False, # Start as closed + max_files_open=sarah_agent.max_files_open, ) # Create new files @@ -8340,13 +8505,16 @@ async def test_attach_files_bulk_mixed_existing_new(server, default_user, sarah_ files_metadata=files_to_attach, visible_content_map=visible_content_map, actor=default_user, + max_files_open=sarah_agent.max_files_open, ) # Should not close any files assert closed_files == [] # Verify all files are now open - open_files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user, is_open_only=True) + open_files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True + ) assert len(open_files) == 3 # Verify existing file was updated @@ -8361,27 +8529,26 @@ async def test_attach_files_bulk_mixed_existing_new(server, default_user, sarah_ async def test_attach_files_bulk_empty_list(server, default_user, sarah_agent): """Test attach_files_bulk with empty file list.""" closed_files = await server.file_agent_manager.attach_files_bulk( - agent_id=sarah_agent.id, - files_metadata=[], - visible_content_map={}, - actor=default_user, + agent_id=sarah_agent.id, files_metadata=[], visible_content_map={}, actor=default_user, max_files_open=sarah_agent.max_files_open ) assert closed_files == [] # Verify no files are attached - attached_files = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user) + attached_files = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user + ) assert len(attached_files) == 0 @pytest.mark.asyncio async def test_attach_files_bulk_oversized_bulk(server, default_user, sarah_agent, default_source): - """Test bulk attach when trying to attach more files than MAX_FILES_OPEN allows.""" - from letta.constants import MAX_FILES_OPEN + """Test bulk attach when trying to attach more files than max_files_open allows.""" + max_files_open = sarah_agent.max_files_open # Create more files than the limit allows oversized_files = [] - for i in range(MAX_FILES_OPEN + 3): # 3 more than limit + for i in range(max_files_open + 3): # 3 more than limit file_metadata = PydanticFileMetadata( file_name=f"oversized_{i}.txt", organization_id=default_user.organization_id, @@ -8390,7 +8557,7 @@ async def test_attach_files_bulk_oversized_bulk(server, default_user, sarah_agen file = await server.file_manager.create_file(file_metadata=file_metadata, actor=default_user, text=f"oversized {i}") oversized_files.append(file) - visible_content_map = {f"oversized_{i}.txt": f"oversized visible {i}" for i in range(MAX_FILES_OPEN + 3)} + visible_content_map = {f"oversized_{i}.txt": f"oversized visible {i}" for i in range(max_files_open + 3)} # Bulk attach all files (more than limit) closed_files = await server.file_agent_manager.attach_files_bulk( @@ -8398,6 +8565,7 @@ async def test_attach_files_bulk_oversized_bulk(server, default_user, sarah_agen files_metadata=oversized_files, visible_content_map=visible_content_map, actor=default_user, + max_files_open=sarah_agent.max_files_open, ) # Should have closed exactly 3 files (the excess) @@ -8406,13 +8574,17 @@ async def test_attach_files_bulk_oversized_bulk(server, default_user, sarah_agen # CRITICAL: Verify no duplicates in closed_files list assert len(closed_files) == len(set(closed_files)), f"Duplicate file names in closed_files: {closed_files}" - # Should have exactly MAX_FILES_OPEN files open - open_files_after = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user, is_open_only=True) - assert len(open_files_after) == MAX_FILES_OPEN + # Should have exactly max_files_open files open + open_files_after = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user, is_open_only=True + ) + assert len(open_files_after) == max_files_open # All files should be attached (some open, some closed) - all_files_after = await server.file_agent_manager.list_files_for_agent(sarah_agent.id, actor=default_user) - assert len(all_files_after) == MAX_FILES_OPEN + 3 + all_files_after = await server.file_agent_manager.list_files_for_agent( + sarah_agent.id, per_file_view_window_char_limit=sarah_agent.per_file_view_window_char_limit, actor=default_user + ) + assert len(all_files_after) == max_files_open + 3 # ====================================================================================================================== diff --git a/tests/test_sources.py b/tests/test_sources.py index 17053cd7..b7e64e2c 100644 --- a/tests/test_sources.py +++ b/tests/test_sources.py @@ -282,7 +282,7 @@ def test_attach_existing_files_creates_source_blocks_correctly(disable_pinecone, - read_only=true - chars_current=46 -- chars_limit=50000 +- chars_limit=15000 [Viewing file start (out of 1 chunks)] @@ -339,7 +339,7 @@ def test_delete_source_removes_source_blocks_correctly(disable_pinecone, client: - read_only=true - chars_current=46 -- chars_limit=50000 +- chars_limit=15000 [Viewing file start (out of 1 chunks)]