feat: Add file status to core memory jinja template (#2604)
This commit is contained in:
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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>"
|
||||
|
||||
@@ -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"""
|
||||
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user