feat: Add file status to core memory jinja template (#2604)

This commit is contained in:
Matthew Zhou
2025-06-03 15:07:21 -07:00
committed by GitHub
parent d0519437af
commit 1baa124c76
6 changed files with 47 additions and 27 deletions

View File

@@ -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,
)

View File

@@ -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"
"</description>\n"
"<metadata>"
"{% 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"
"</metadata>\n"
"<value>\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</memory_blocks>"
"<files>\nThe following memory blocks are currently accessible in your core memory unit:\n\n"
"<files>\nThe following memory files are currently accessible:\n\n"
"{% for block in file_blocks %}"
f"<file status=\"{{{{ '{FileStatus.open.value}' if block.value else '{FileStatus.closed.value}' }}}}\">\n"
"<{{ block.label }}>\n"
"<description>\n"
"{{ block.description }}\n"
"</description>\n"
"<metadata>"
"{% 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"
"</metadata>\n"
"<value>\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 %}"
"</value>\n"
"</{{ block.label }}>\n"
"</file>\n"
"{% if not loop.last %}\n{% endif %}"
"{% endfor %}"
"\n</memory_blocks>"
"</files>"
"\n</files>"
)
# 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"
"</description>\n"
"<metadata>"
"{% 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"
"</metadata>\n"
"<value>\n"
"{{ block.value }}\n"
@@ -364,18 +373,22 @@ def get_prompt_template_for_agent_type(agent_type: Optional[AgentType] = None):
"{% endfor %}"
"\n</memory_blocks>"
"<files>\nThe following memory files are currently accessible:\n\n"
"{% for block in file_blocks%}"
"{% for block in file_blocks %}"
f"<file status=\"{{{{ '{FileStatus.open.value}' if block.value else '{FileStatus.closed.value}' }}}}\">\n"
"<{{ block.label }}>\n"
"<description>\n"
"{{ block.description }}\n"
"</description>\n"
"<metadata>"
"{% 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"
"</metadata>\n"
"<value>\n"
"{{ block.value }}\n"
"</value>\n"
"</{{ block.label }}>\n"
"</file>\n"
"{% if not loop.last %}\n{% endif %}"
"{% endfor %}"
"\n</files>"

View File

@@ -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"""

View File

@@ -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":

View File

@@ -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 "<file>" in overview.core_memory
assert '<file status="open">' in overview.core_memory
assert "</file>" in overview.core_memory

View File

@@ -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