From 1baa124c769a2b4d533c4702ee0e71b62e5a9cb8 Mon Sep 17 00:00:00 2001 From: Matthew Zhou Date: Tue, 3 Jun 2025 15:07:21 -0700 Subject: [PATCH] feat: Add file status to core memory jinja template (#2604) --- letta/orm/files_agents.py | 18 ++++++++---------- letta/schemas/agent.py | 29 +++++++++++++++++++++-------- letta/schemas/file.py | 10 ++++++++++ letta/services/agent_manager.py | 2 +- tests/helpers/utils.py | 2 +- tests/test_managers.py | 13 ++++++------- 6 files changed, 47 insertions(+), 27 deletions(-) diff --git a/letta/orm/files_agents.py b/letta/orm/files_agents.py index ac9d9e34..19b25bca 100644 --- a/letta/orm/files_agents.py +++ b/letta/orm/files_agents.py @@ -58,13 +58,11 @@ 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) -> Optional[PydanticBlock]: - if self.is_open: - return PydanticBlock( - organization_id=self.organization_id, - value=self.visible_content if self.visible_content else "", - label=self.file.file_name, - read_only=True, - ) - else: - return None + def to_pydantic_block(self) -> PydanticBlock: + visible_content = self.visible_content if self.visible_content and self.is_open else "" + return PydanticBlock( + organization_id=self.organization_id, + value=visible_content, + label=self.file.file_name, + read_only=True, + ) diff --git a/letta/schemas/agent.py b/letta/schemas/agent.py index b94fd93a..32e9cf33 100644 --- a/letta/schemas/agent.py +++ b/letta/schemas/agent.py @@ -8,6 +8,7 @@ from letta.helpers import ToolRulesSolver from letta.schemas.block import CreateBlock from letta.schemas.embedding_config import EmbeddingConfig from letta.schemas.environment_variables import AgentEnvironmentVariable +from letta.schemas.file import FileStatus from letta.schemas.group import Group from letta.schemas.letta_base import OrmMetadataBase from letta.schemas.llm_config import LLMConfig @@ -311,7 +312,9 @@ def get_prompt_template_for_agent_type(agent_type: Optional[AgentType] = None): "{{ block.description }}\n" "\n" "" - "{% if block.read_only %}\n- read_only=true{% endif %}\n- chars_current={{ block.value|length }}\n- chars_limit={{ block.limit }}\n" + "{% if block.read_only %}\n- read_only=true{% endif %}\n" + "- chars_current={{ block.value|length }}\n" + "- chars_limit={{ block.limit }}\n" "\n" "\n" f"{CORE_MEMORY_LINE_NUMBER_WARNING}\n" @@ -323,14 +326,17 @@ def get_prompt_template_for_agent_type(agent_type: Optional[AgentType] = None): "{% if not loop.last %}\n{% endif %}" "{% endfor %}" "\n" - "\nThe following memory blocks are currently accessible in your core memory unit:\n\n" + "\nThe following memory files are currently accessible:\n\n" "{% for block in file_blocks %}" + f"\n" "<{{ block.label }}>\n" "\n" "{{ block.description }}\n" "\n" "" - "{% if block.read_only %}\n- read_only=true{% endif %}\n- chars_current={{ block.value|length }}\n- chars_limit={{ block.limit }}\n" + "{% if block.read_only %}\n- read_only=true{% endif %}\n" + "- chars_current={{ block.value|length }}\n" + "- chars_limit={{ block.limit }}\n" "\n" "\n" f"{CORE_MEMORY_LINE_NUMBER_WARNING}\n" @@ -339,11 +345,12 @@ def get_prompt_template_for_agent_type(agent_type: Optional[AgentType] = None): "{% endfor %}" "\n" "\n" + "\n" "{% if not loop.last %}\n{% endif %}" "{% endfor %}" - "\n" - "" + "\n" ) + # Default setup (MemGPT), no line numbers else: return ( @@ -354,7 +361,9 @@ def get_prompt_template_for_agent_type(agent_type: Optional[AgentType] = None): "{{ block.description }}\n" "\n" "" - "{% if block.read_only %}\n- read_only=true{% endif %}\n- chars_current={{ block.value|length }}\n- chars_limit={{ block.limit }}\n" + "{% if block.read_only %}\n- read_only=true{% endif %}\n" + "- chars_current={{ block.value|length }}\n" + "- chars_limit={{ block.limit }}\n" "\n" "\n" "{{ block.value }}\n" @@ -364,18 +373,22 @@ def get_prompt_template_for_agent_type(agent_type: Optional[AgentType] = None): "{% endfor %}" "\n" "\nThe following memory files are currently accessible:\n\n" - "{% for block in file_blocks%}" + "{% for block in file_blocks %}" + f"\n" "<{{ block.label }}>\n" "\n" "{{ block.description }}\n" "\n" "" - "{% if block.read_only %}\n- read_only=true{% endif %}\n- chars_current={{ block.value|length }}\n- chars_limit={{ block.limit }}\n" + "{% if block.read_only %}\n- read_only=true{% endif %}\n" + "- chars_current={{ block.value|length }}\n" + "- chars_limit={{ block.limit }}\n" "\n" "\n" "{{ block.value }}\n" "\n" "\n" + "\n" "{% if not loop.last %}\n{% endif %}" "{% endfor %}" "\n" diff --git a/letta/schemas/file.py b/letta/schemas/file.py index f537485d..ecd9fe0a 100644 --- a/letta/schemas/file.py +++ b/letta/schemas/file.py @@ -1,4 +1,5 @@ from datetime import datetime +from enum import Enum from typing import Optional from pydantic import Field @@ -6,6 +7,15 @@ from pydantic import Field from letta.schemas.letta_base import LettaBase +class FileStatus(str, Enum): + """ + Enum to represent the state of a file. + """ + + open = "open" + closed = "closed" + + class FileMetadataBase(LettaBase): """Base class for FileMetadata schemas""" diff --git a/letta/services/agent_manager.py b/letta/services/agent_manager.py index 2808b03a..1a7b60ab 100644 --- a/letta/services/agent_manager.py +++ b/letta/services/agent_manager.py @@ -2622,7 +2622,7 @@ class AgentManager: return results async def get_context_window(self, agent_id: str, actor: PydanticUser) -> ContextWindowOverview: - agent_state = await self.get_agent_by_id_async(agent_id=agent_id, actor=actor) + agent_state = await self.rebuild_system_prompt_async(agent_id=agent_id, actor=actor, force=True) calculator = ContextWindowCalculator() if os.getenv("LETTA_ENVIRONMENT") == "PRODUCTION" or agent_state.llm_config.model_endpoint_type == "anthropic": diff --git a/tests/helpers/utils.py b/tests/helpers/utils.py index 7674ee1a..2dfa5cca 100644 --- a/tests/helpers/utils.py +++ b/tests/helpers/utils.py @@ -236,5 +236,5 @@ def validate_context_window_overview(overview: ContextWindowOverview, attached_f # 16. Check attached file is visible if attached_file: assert attached_file.visible_content in overview.core_memory - assert "" in overview.core_memory + assert '' in overview.core_memory assert "" in overview.core_memory diff --git a/tests/test_managers.py b/tests/test_managers.py index 2f2eee28..f8a42d82 100644 --- a/tests/test_managers.py +++ b/tests/test_managers.py @@ -696,7 +696,7 @@ async def test_get_context_window_basic(server: SyncServer, comprehensive_test_a comprehensive_agent_checks(created_agent, create_agent_request, actor=default_user) # Attach a file - await server.file_agent_manager.attach_file( + assoc = await server.file_agent_manager.attach_file( agent_id=created_agent.id, file_id=default_file.id, actor=default_user, @@ -705,7 +705,7 @@ async def test_get_context_window_basic(server: SyncServer, comprehensive_test_a # Get context window and check for basic appearances context_window_overview = await server.agent_manager.get_context_window(agent_id=created_agent.id, actor=default_user) - validate_context_window_overview(context_window_overview) + validate_context_window_overview(context_window_overview, assoc) # Test deleting the agent server.agent_manager.delete_agent(created_agent.id, default_user) @@ -5851,7 +5851,9 @@ async def test_attach_is_idempotent(server, default_user, sarah_agent, default_f sarah_agent = await server.agent_manager.get_agent_by_id_async(agent_id=sarah_agent.id, actor=default_user) file_blocks = sarah_agent.memory.file_blocks - assert len(file_blocks) == 0 # Is not open + assert len(file_blocks) == 1 + assert file_blocks[0].value == "" # not open + assert file_blocks[0].label == default_file.file_name @pytest.mark.asyncio @@ -5915,10 +5917,7 @@ async def test_list_files_and_agents( sarah_agent = await server.agent_manager.get_agent_by_id_async(agent_id=sarah_agent.id, actor=default_user) file_blocks = sarah_agent.memory.file_blocks - assert len(file_blocks) == 1 - assert file_blocks[0].value == "" - assert file_blocks[0].label == default_file.file_name - + assert len(file_blocks) == 2 charles_agent = await server.agent_manager.get_agent_by_id_async(agent_id=charles_agent.id, actor=default_user) file_blocks = charles_agent.memory.file_blocks assert len(file_blocks) == 1