From 99c479cd787e7d7b078034fb830aa2bca3777502 Mon Sep 17 00:00:00 2001 From: Kevin Lin Date: Wed, 23 Jul 2025 17:05:53 -0700 Subject: [PATCH] feat: add max files and file open window to system prompt (#3515) --- letta/agents/base_agent.py | 5 ++++- letta/agents/voice_agent.py | 1 + letta/schemas/agent.py | 18 ++++++++++++++++++ letta/schemas/memory.py | 12 +++++++++--- letta/services/agent_manager.py | 12 ++++++++++-- letta/services/helpers/agent_manager_helper.py | 6 +++++- tests/helpers/utils.py | 2 ++ 7 files changed, 49 insertions(+), 7 deletions(-) diff --git a/letta/agents/base_agent.py b/letta/agents/base_agent.py index a556ab73..1e5401b2 100644 --- a/letta/agents/base_agent.py +++ b/letta/agents/base_agent.py @@ -122,7 +122,9 @@ class BaseAgent(ABC): curr_dynamic_section = extract_dynamic_section(curr_system_message_text) # generate just the memory string with current state for comparison - curr_memory_str = agent_state.memory.compile(tool_usage_rules=tool_constraint_block, sources=agent_state.sources) + curr_memory_str = agent_state.memory.compile( + tool_usage_rules=tool_constraint_block, sources=agent_state.sources, max_files_open=agent_state.max_files_open + ) new_dynamic_section = extract_dynamic_section(curr_memory_str) # compare just the dynamic sections (memory blocks, tool rules, directories) @@ -149,6 +151,7 @@ class BaseAgent(ABC): archival_memory_size=num_archival_memories, tool_rules_solver=tool_rules_solver, sources=agent_state.sources, + max_files_open=agent_state.max_files_open, ) diff = united_diff(curr_system_message_text, new_system_message_str) diff --git a/letta/agents/voice_agent.py b/letta/agents/voice_agent.py index 9e30cb54..e28fb17f 100644 --- a/letta/agents/voice_agent.py +++ b/letta/agents/voice_agent.py @@ -153,6 +153,7 @@ class VoiceAgent(BaseAgent): previous_message_count=self.num_messages, archival_memory_size=self.num_archival_memories, sources=agent_state.sources, + max_files_open=agent_state.max_files_open, ) letta_message_db_queue = create_input_messages( input_messages=input_messages, agent_id=agent_state.id, timezone=agent_state.timezone, actor=self.actor diff --git a/letta/schemas/agent.py b/letta/schemas/agent.py index 54f5796d..bdc36083 100644 --- a/letta/schemas/agent.py +++ b/letta/schemas/agent.py @@ -360,6 +360,12 @@ def get_prompt_template_for_agent_type(agent_type: Optional[AgentType] = None): return ( "{% if sources %}" "\n" + "{% if max_files_open %}" + "\n" + "- current_files_open={{ file_blocks|selectattr('value')|list|length }}\n" + "- max_files_open={{ max_files_open }}\n" + "\n" + "{% endif %}" "{% for source in sources %}" f'\n' "{% if source.description %}" @@ -427,6 +433,12 @@ def get_prompt_template_for_agent_type(agent_type: Optional[AgentType] = None): "{% endif %}" "\n\n{% if sources %}" "\n" + "{% if max_files_open %}" + "\n" + "- current_files_open={{ file_blocks|selectattr('value')|list|length }}\n" + "- max_files_open={{ max_files_open }}\n" + "\n" + "{% endif %}" "{% for source in sources %}" f'\n' "{% if source.description %}" @@ -493,6 +505,12 @@ def get_prompt_template_for_agent_type(agent_type: Optional[AgentType] = None): "{% endif %}" "\n\n{% if sources %}" "\n" + "{% if max_files_open %}" + "\n" + "- current_files_open={{ file_blocks|selectattr('value')|list|length }}\n" + "- max_files_open={{ max_files_open }}\n" + "\n" + "{% endif %}" "{% for source in sources %}" f'\n' "{% if source.description %}" diff --git a/letta/schemas/memory.py b/letta/schemas/memory.py index bb856676..c240bded 100644 --- a/letta/schemas/memory.py +++ b/letta/schemas/memory.py @@ -124,7 +124,7 @@ class Memory(BaseModel, validate_assignment=True): Template(prompt_template) # Validate compatibility with current memory structure - Template(prompt_template).render(blocks=self.blocks, file_blocks=self.file_blocks, sources=[]) + Template(prompt_template).render(blocks=self.blocks, file_blocks=self.file_blocks, sources=[], max_files_open=None) # If we get here, the template is valid and compatible self.prompt_template = prompt_template @@ -133,11 +133,17 @@ class Memory(BaseModel, validate_assignment=True): except Exception as e: raise ValueError(f"Prompt template is not compatible with current memory structure: {str(e)}") - def compile(self, tool_usage_rules=None, sources=None) -> str: + def compile(self, tool_usage_rules=None, sources=None, max_files_open=None) -> str: """Generate a string representation of the memory in-context using the Jinja2 template""" try: template = Template(self.prompt_template) - return template.render(blocks=self.blocks, file_blocks=self.file_blocks, tool_usage_rules=tool_usage_rules, sources=sources) + return template.render( + blocks=self.blocks, + file_blocks=self.file_blocks, + tool_usage_rules=tool_usage_rules, + sources=sources, + max_files_open=max_files_open, + ) except TemplateSyntaxError as e: raise ValueError(f"Invalid Jinja2 template syntax: {str(e)}") except Exception as e: diff --git a/letta/services/agent_manager.py b/letta/services/agent_manager.py index b7b02909..6302f645 100644 --- a/letta/services/agent_manager.py +++ b/letta/services/agent_manager.py @@ -1587,6 +1587,7 @@ class AgentManager: previous_message_count=num_messages - len(agent_state.message_ids), archival_memory_size=num_archival_memories, sources=agent_state.sources, + max_files_open=agent_state.max_files_open, ) diff = united_diff(curr_system_message_openai["content"], new_system_message_str) @@ -1644,7 +1645,9 @@ class AgentManager: # note: we only update the system prompt if the core memory is changed # this means that the archival/recall memory statistics may be someout out of date curr_memory_str = agent_state.memory.compile( - sources=agent_state.sources, tool_usage_rules=tool_rules_solver.compile_tool_rule_prompts() + sources=agent_state.sources, + tool_usage_rules=tool_rules_solver.compile_tool_rule_prompts(), + max_files_open=agent_state.max_files_open, ) if curr_memory_str in curr_system_message_openai["content"] and not force: # NOTE: could this cause issues if a block is removed? (substring match would still work) @@ -1672,6 +1675,7 @@ class AgentManager: archival_memory_size=num_archival_memories, tool_rules_solver=tool_rules_solver, sources=agent_state.sources, + max_files_open=agent_state.max_files_open, ) diff = united_diff(curr_system_message_openai["content"], new_system_message_str) @@ -1831,7 +1835,11 @@ class AgentManager: system_message = await self.message_manager.get_message_by_id_async(message_id=agent_state.message_ids[0], actor=actor) temp_tool_rules_solver = ToolRulesSolver(agent_state.tool_rules) if ( - new_memory.compile(sources=agent_state.sources, tool_usage_rules=temp_tool_rules_solver.compile_tool_rule_prompts()) + new_memory.compile( + sources=agent_state.sources, + tool_usage_rules=temp_tool_rules_solver.compile_tool_rule_prompts(), + max_files_open=agent_state.max_files_open, + ) not in system_message.content[0].text ): # update the blocks (LRW) in the DB diff --git a/letta/services/helpers/agent_manager_helper.py b/letta/services/helpers/agent_manager_helper.py index e5ff86a4..ae13435f 100644 --- a/letta/services/helpers/agent_manager_helper.py +++ b/letta/services/helpers/agent_manager_helper.py @@ -259,6 +259,7 @@ def compile_system_message( archival_memory_size: int = 0, tool_rules_solver: Optional[ToolRulesSolver] = None, sources: Optional[List] = None, + max_files_open: Optional[int] = None, ) -> str: """Prepare the final/full system message that will be fed into the LLM API @@ -291,7 +292,9 @@ def compile_system_message( timezone=timezone, ) - memory_with_sources = in_context_memory.compile(tool_usage_rules=tool_constraint_block, sources=sources) + memory_with_sources = in_context_memory.compile( + tool_usage_rules=tool_constraint_block, sources=sources, max_files_open=max_files_open + ) full_memory_string = memory_with_sources + "\n\n" + memory_metadata_string # Add to the variables list to inject @@ -343,6 +346,7 @@ def initialize_message_sequence( previous_message_count=previous_message_count, archival_memory_size=archival_memory_size, sources=agent_state.sources, + max_files_open=agent_state.max_files_open, ) first_user_message = get_login_event(agent_state.timezone) # event letting Letta know the user just logged in diff --git a/tests/helpers/utils.py b/tests/helpers/utils.py index 950e44d0..3d74e430 100644 --- a/tests/helpers/utils.py +++ b/tests/helpers/utils.py @@ -240,6 +240,8 @@ def validate_context_window_overview( assert attached_file.visible_content in overview.core_memory, "File must be attached in core memory" assert '" in overview.core_memory + assert "max_files_open" in overview.core_memory, "Max files should be set in core memory" + assert "current_files_open" in overview.core_memory, "Current files should be set in core memory" # Check for tools assert overview.num_tokens_functions_definitions > 0