chore: Rename agent file manager (#3496)

This commit is contained in:
Matthew Zhou
2025-07-22 17:40:12 -07:00
committed by GitHub
parent 2744a93375
commit 6020632afd
5 changed files with 143 additions and 740 deletions

View File

@@ -79,6 +79,7 @@ from letta.server.rest_api.chat_completions_interface import ChatCompletionsStre
from letta.server.rest_api.interface import StreamingServerInterface
from letta.server.rest_api.utils import sse_async_generator
from letta.services.agent_manager import AgentManager
from letta.services.agent_serialization_manager import AgentFileManager
from letta.services.block_manager import BlockManager
from letta.services.file_manager import FileManager
from letta.services.file_processor.chunker.line_chunker import LineChunker
@@ -223,6 +224,7 @@ class SyncServer(Server):
self.telemetry_manager = TelemetryManager()
self.file_agent_manager = FileAgentManager()
self.file_manager = FileManager()
self.agent_serialization_manager = AgentFileManager()
# A resusable httpx client
timeout = httpx.Timeout(connect=10.0, read=20.0, write=10.0, pool=10.0)

737
poetry.lock generated

File diff suppressed because it is too large Load Diff

View File

@@ -72,7 +72,7 @@ llama-index = "^0.12.2"
llama-index-embeddings-openai = "^0.3.1"
e2b-code-interpreter = {version = "^1.0.3", optional = true}
anthropic = "^0.49.0"
letta_client = "^0.1.197"
letta_client = "^0.1.213"
openai = "^1.60.0"
opentelemetry-api = "1.30.0"
opentelemetry-sdk = "1.30.0"

View File

@@ -26,7 +26,7 @@ from letta.schemas.organization import Organization
from letta.schemas.source import Source
from letta.schemas.user import User
from letta.server.server import SyncServer
from letta.services.agent_file_manager import AgentFileManager
from letta.services.agent_serialization_manager import AgentFileManager
from letta.services.file_processor.embedder.openai_embedder import OpenAIEmbedder
from letta.services.file_processor.parser.markitdown_parser import MarkitdownFileParser
from letta.services.file_processor.parser.mistral_parser import MistralFileParser
@@ -158,7 +158,7 @@ def test_block(server: SyncServer, default_user):
@pytest.fixture
def agent_file_manager(server, default_user):
def agent_serialization_manager(server, default_user):
"""Fixture to create AgentFileManager with all required services including file processing."""
manager = AgentFileManager(
agent_manager=server.agent_manager,
@@ -674,11 +674,11 @@ def validate_id_format(schema: AgentFileSchema) -> bool:
class TestFileExport:
"""Test file export functionality with comprehensive validation"""
async def test_basic_file_export(self, default_user, agent_file_manager, agent_with_files):
async def test_basic_file_export(self, default_user, agent_serialization_manager, agent_with_files):
"""Test basic file export functionality"""
agent_id, source_id, file_id = agent_with_files
exported = await agent_file_manager.export([agent_id], actor=default_user)
exported = await agent_serialization_manager.export([agent_id], actor=default_user)
assert len(exported.agents) == 1
assert len(exported.sources) == 1
@@ -696,7 +696,7 @@ class TestFileExport:
assert file_agent.file_id == exported.files[0].id
assert file_agent.source_id == exported.sources[0].id
async def test_multiple_files_per_source(self, server, default_user, agent_file_manager):
async def test_multiple_files_per_source(self, server, default_user, agent_serialization_manager):
"""Test export with multiple files from the same source"""
source = await create_test_source(server, "multi-file-source", default_user)
file1 = await create_test_file(server, "file1.txt", source.id, default_user)
@@ -704,7 +704,7 @@ class TestFileExport:
agent = await create_test_agent_with_files(server, "multi-file-agent", default_user, [(source.id, file1.id), (source.id, file2.id)])
exported = await agent_file_manager.export([agent.id], actor=default_user)
exported = await agent_serialization_manager.export([agent.id], actor=default_user)
assert len(exported.agents) == 1
assert len(exported.sources) == 1
@@ -722,7 +722,7 @@ class TestFileExport:
assert file_agent.file_id in file_ids
assert file_agent.source_id == source_id
async def test_multiple_sources_export(self, server, default_user, agent_file_manager):
async def test_multiple_sources_export(self, server, default_user, agent_serialization_manager):
"""Test export with files from multiple sources"""
source1 = await create_test_source(server, "source-1", default_user)
source2 = await create_test_source(server, "source-2", default_user)
@@ -733,7 +733,7 @@ class TestFileExport:
server, "multi-source-agent", default_user, [(source1.id, file1.id), (source2.id, file2.id)]
)
exported = await agent_file_manager.export([agent.id], actor=default_user)
exported = await agent_serialization_manager.export([agent.id], actor=default_user)
assert len(exported.agents) == 1
assert len(exported.sources) == 2
@@ -743,7 +743,7 @@ class TestFileExport:
for file_schema in exported.files:
assert file_schema.source_id in source_ids
async def test_cross_agent_file_deduplication(self, server, default_user, agent_file_manager):
async def test_cross_agent_file_deduplication(self, server, default_user, agent_serialization_manager):
"""Test that files shared across agents are deduplicated in export"""
source = await create_test_source(server, "shared-source", default_user)
shared_file = await create_test_file(server, "shared.txt", source.id, default_user)
@@ -751,7 +751,7 @@ class TestFileExport:
agent1 = await create_test_agent_with_files(server, "agent-1", default_user, [(source.id, shared_file.id)])
agent2 = await create_test_agent_with_files(server, "agent-2", default_user, [(source.id, shared_file.id)])
exported = await agent_file_manager.export([agent1.id, agent2.id], actor=default_user)
exported = await agent_serialization_manager.export([agent1.id, agent2.id], actor=default_user)
assert len(exported.agents) == 2
assert len(exported.sources) == 1
@@ -766,14 +766,14 @@ class TestFileExport:
assert file_agent.file_id == file_id
assert file_agent.source_id == source_id
async def test_file_agent_relationship_preservation(self, server, default_user, agent_file_manager):
async def test_file_agent_relationship_preservation(self, server, default_user, agent_serialization_manager):
"""Test that file-agent relationship details are preserved"""
source = await create_test_source(server, "test-source", default_user)
file = await create_test_file(server, "test.txt", source.id, default_user)
agent = await create_test_agent_with_files(server, "test-agent", default_user, [(source.id, file.id)])
exported = await agent_file_manager.export([agent.id], actor=default_user)
exported = await agent_serialization_manager.export([agent.id], actor=default_user)
agent = exported.agents[0]
file_agent = agent.files_agents[0]
@@ -782,13 +782,13 @@ class TestFileExport:
assert file_agent.is_open is True
assert hasattr(file_agent, "last_accessed_at")
async def test_id_remapping_consistency(self, server, default_user, agent_file_manager):
async def test_id_remapping_consistency(self, server, default_user, agent_serialization_manager):
"""Test that ID remapping is consistent across all references"""
source = await create_test_source(server, "consistency-source", default_user)
file = await create_test_file(server, "consistency.txt", source.id, default_user)
agent = await create_test_agent_with_files(server, "consistency-agent", default_user, [(source.id, file.id)])
exported = await agent_file_manager.export([agent.id], actor=default_user)
exported = await agent_serialization_manager.export([agent.id], actor=default_user)
agent_schema = exported.agents[0]
source_schema = exported.sources[0]
@@ -800,7 +800,7 @@ class TestFileExport:
assert file_agent.file_id == file_schema.id
assert file_agent.source_id == source_schema.id
async def test_empty_file_relationships(self, server, default_user, agent_file_manager):
async def test_empty_file_relationships(self, server, default_user, agent_serialization_manager):
"""Test export of agent with no file relationships"""
agent_create = CreateAgent(
name="no-files-agent",
@@ -809,7 +809,7 @@ class TestFileExport:
)
agent = await server.agent_manager.create_agent_async(agent_create, actor=default_user)
exported = await agent_file_manager.export([agent.id], actor=default_user)
exported = await agent_serialization_manager.export([agent.id], actor=default_user)
assert len(exported.agents) == 1
assert len(exported.sources) == 0
@@ -818,11 +818,11 @@ class TestFileExport:
agent_schema = exported.agents[0]
assert len(agent_schema.files_agents) == 0
async def test_file_content_inclusion_in_export(self, default_user, agent_file_manager, agent_with_files):
async def test_file_content_inclusion_in_export(self, default_user, agent_serialization_manager, agent_with_files):
"""Test that file content is included in export"""
agent_id, source_id, file_id = agent_with_files
exported = await agent_file_manager.export([agent_id], actor=default_user)
exported = await agent_serialization_manager.export([agent_id], actor=default_user)
file_schema = exported.files[0]
assert hasattr(file_schema, "content") or file_schema.content is not None
@@ -831,9 +831,9 @@ class TestFileExport:
class TestAgentFileExport:
"""Tests for agent file export functionality."""
async def test_basic_export(self, agent_file_manager, test_agent, default_user):
async def test_basic_export(self, agent_serialization_manager, test_agent, default_user):
"""Test basic agent export functionality."""
agent_file = await agent_file_manager.export([test_agent.id], default_user)
agent_file = await agent_serialization_manager.export([test_agent.id], default_user)
assert isinstance(agent_file, AgentFileSchema)
assert len(agent_file.agents) == 1
@@ -852,7 +852,7 @@ class TestAgentFileExport:
assert len(exported_agent.messages) > 0
assert len(exported_agent.in_context_message_ids) > 0
async def test_export_multiple_agents(self, server, agent_file_manager, test_agent, default_user, weather_tool):
async def test_export_multiple_agents(self, server, agent_serialization_manager, test_agent, default_user, weather_tool):
"""Test exporting multiple agents."""
create_agent_request = CreateAgent(
name="second_test_agent",
@@ -870,7 +870,7 @@ class TestAgentFileExport:
actor=default_user,
)
agent_file = await agent_file_manager.export([test_agent.id, second_agent.id], default_user)
agent_file = await agent_serialization_manager.export([test_agent.id, second_agent.id], default_user)
assert len(agent_file.agents) == 2
assert validate_id_format(agent_file)
@@ -878,9 +878,9 @@ class TestAgentFileExport:
agent_ids = {agent.id for agent in agent_file.agents}
assert len(agent_ids) == 2
async def test_export_id_remapping(self, agent_file_manager, test_agent, default_user):
async def test_export_id_remapping(self, agent_serialization_manager, test_agent, default_user):
"""Test that IDs are properly remapped during export."""
agent_file = await agent_file_manager.export([test_agent.id], default_user)
agent_file = await agent_serialization_manager.export([test_agent.id], default_user)
exported_agent = agent_file.agents[0]
@@ -899,9 +899,9 @@ class TestAgentFileExport:
for in_context_id in exported_agent.in_context_message_ids:
assert in_context_id in message_ids, f"In-context message ID {in_context_id} not found in messages"
async def test_message_agent_id_remapping(self, agent_file_manager, test_agent, default_user):
async def test_message_agent_id_remapping(self, agent_serialization_manager, test_agent, default_user):
"""Test that message.agent_id is properly remapped during export."""
agent_file = await agent_file_manager.export([test_agent.id], default_user)
agent_file = await agent_serialization_manager.export([test_agent.id], default_user)
exported_agent = agent_file.agents[0]
@@ -913,22 +913,22 @@ class TestAgentFileExport:
assert exported_agent.id == "agent-0"
assert exported_agent.id != test_agent.id
async def test_export_empty_agent_list(self, agent_file_manager, default_user):
async def test_export_empty_agent_list(self, agent_serialization_manager, default_user):
"""Test exporting empty agent list."""
agent_file = await agent_file_manager.export([], default_user)
agent_file = await agent_serialization_manager.export([], default_user)
assert len(agent_file.agents) == 0
assert len(agent_file.tools) == 0
assert len(agent_file.blocks) == 0
async def test_export_nonexistent_agent(self, agent_file_manager, default_user):
async def test_export_nonexistent_agent(self, agent_serialization_manager, default_user):
"""Test exporting non-existent agent raises error."""
with pytest.raises(AgentFileExportError): # Should raise AgentFileExportError for non-existent agent
await agent_file_manager.export(["non-existent-id"], default_user)
await agent_serialization_manager.export(["non-existent-id"], default_user)
async def test_revision_id_automatic_setting(self, agent_file_manager, test_agent, default_user):
async def test_revision_id_automatic_setting(self, agent_serialization_manager, test_agent, default_user):
"""Test that revision_id is automatically set to the latest alembic revision."""
agent_file = await agent_file_manager.export([test_agent.id], default_user)
agent_file = await agent_serialization_manager.export([test_agent.id], default_user)
from letta.utils import get_latest_alembic_revision
@@ -945,11 +945,11 @@ class TestAgentFileExport:
class TestAgentFileImport:
"""Tests for agent file import functionality."""
async def test_basic_import(self, agent_file_manager, test_agent, default_user, other_user):
async def test_basic_import(self, agent_serialization_manager, test_agent, default_user, other_user):
"""Test basic agent import functionality."""
agent_file = await agent_file_manager.export([test_agent.id], default_user)
agent_file = await agent_serialization_manager.export([test_agent.id], default_user)
result = await agent_file_manager.import_file(agent_file, other_user)
result = await agent_serialization_manager.import_file(agent_file, other_user)
assert result.success
assert result.imported_count > 0
@@ -959,11 +959,11 @@ class TestAgentFileImport:
if file_id.startswith("agent-"):
assert db_id != test_agent.id # New agent should have different ID
async def test_import_preserves_data(self, server, agent_file_manager, test_agent, default_user, other_user):
async def test_import_preserves_data(self, server, agent_serialization_manager, test_agent, default_user, other_user):
"""Test that import preserves all important data."""
agent_file = await agent_file_manager.export([test_agent.id], default_user)
agent_file = await agent_serialization_manager.export([test_agent.id], default_user)
result = await agent_file_manager.import_file(agent_file, other_user)
result = await agent_serialization_manager.import_file(agent_file, other_user)
imported_agent_id = next(db_id for file_id, db_id in result.id_mappings.items() if file_id == "agent-0")
imported_agent = server.agent_manager.get_agent_by_id(imported_agent_id, other_user)
@@ -987,11 +987,11 @@ class TestAgentFileImport:
assert orig_msg.content == imp_msg.content
assert imp_msg.agent_id == imported_agent_id # Should be remapped to new agent
async def test_import_message_context_preservation(self, server, agent_file_manager, test_agent, default_user, other_user):
async def test_import_message_context_preservation(self, server, agent_serialization_manager, test_agent, default_user, other_user):
"""Test that in-context message references are preserved during import."""
agent_file = await agent_file_manager.export([test_agent.id], default_user)
agent_file = await agent_serialization_manager.export([test_agent.id], default_user)
result = await agent_file_manager.import_file(agent_file, other_user)
result = await agent_serialization_manager.import_file(agent_file, other_user)
imported_agent_id = next(db_id for file_id, db_id in result.id_mappings.items() if file_id == "agent-0")
imported_agent = server.agent_manager.get_agent_by_id(imported_agent_id, other_user)
@@ -1004,18 +1004,18 @@ class TestAgentFileImport:
for in_context_id in imported_agent.message_ids:
assert in_context_id in imported_message_ids
async def test_dry_run_import(self, agent_file_manager, test_agent, default_user, other_user):
async def test_dry_run_import(self, agent_serialization_manager, test_agent, default_user, other_user):
"""Test dry run import validation."""
agent_file = await agent_file_manager.export([test_agent.id], default_user)
agent_file = await agent_serialization_manager.export([test_agent.id], default_user)
result = await agent_file_manager.import_file(agent_file, other_user, dry_run=True)
result = await agent_serialization_manager.import_file(agent_file, other_user, dry_run=True)
assert result.success
assert result.imported_count == 0 # No actual imports in dry run
assert len(result.id_mappings) == 0
assert "dry run" in result.message.lower()
async def test_import_validation_errors(self, agent_file_manager, other_user):
async def test_import_validation_errors(self, agent_serialization_manager, other_user):
"""Test import validation catches errors."""
from letta.utils import get_latest_alembic_revision
@@ -1035,13 +1035,13 @@ class TestAgentFileImport:
)
with pytest.raises(AgentFileImportError):
await agent_file_manager.import_file(invalid_agent_file, other_user)
await agent_serialization_manager.import_file(invalid_agent_file, other_user)
class TestAgentFileImportWithProcessing:
"""Tests for agent file import with file processing (chunking/embedding)."""
async def test_import_with_file_processing(self, server, agent_file_manager, default_user, other_user):
async def test_import_with_file_processing(self, server, agent_serialization_manager, default_user, other_user):
"""Test that import processes files for chunking and embedding."""
source = await create_test_source(server, "processing-source", default_user)
file_content = "This is test content for processing. It should be chunked and embedded during import."
@@ -1049,9 +1049,9 @@ class TestAgentFileImportWithProcessing:
agent = await create_test_agent_with_files(server, "processing-agent", default_user, [(source.id, file_metadata.id)])
exported = await agent_file_manager.export([agent.id], default_user)
exported = await agent_serialization_manager.export([agent.id], default_user)
result = await agent_file_manager.import_file(exported, other_user)
result = await agent_serialization_manager.import_file(exported, other_user)
assert result.success
assert result.imported_count > 0
@@ -1061,7 +1061,7 @@ class TestAgentFileImportWithProcessing:
imported_file = await server.file_manager.get_file_by_id(imported_file_id, other_user)
assert imported_file.processing_status.value == "completed"
async def test_import_passage_creation(self, server, agent_file_manager, default_user, other_user):
async def test_import_passage_creation(self, server, agent_serialization_manager, default_user, other_user):
"""Test that import creates passages for file content."""
source = await create_test_source(server, "passage-source", default_user)
file_content = "This content should create passages. Each sentence should be chunked separately."
@@ -1069,9 +1069,9 @@ class TestAgentFileImportWithProcessing:
agent = await create_test_agent_with_files(server, "passage-agent", default_user, [(source.id, file_metadata.id)])
exported = await agent_file_manager.export([agent.id], default_user)
exported = await agent_serialization_manager.export([agent.id], default_user)
result = await agent_file_manager.import_file(exported, other_user)
result = await agent_serialization_manager.import_file(exported, other_user)
imported_file_id = next(db_id for file_id, db_id in result.id_mappings.items() if file_id.startswith("file-"))
@@ -1082,16 +1082,16 @@ class TestAgentFileImportWithProcessing:
assert passage.embedding is not None
assert len(passage.embedding) > 0
async def test_import_file_status_updates(self, server, agent_file_manager, default_user, other_user):
async def test_import_file_status_updates(self, server, agent_serialization_manager, default_user, other_user):
"""Test that file processing status is updated correctly during import."""
source = await create_test_source(server, "status-source", default_user)
file_metadata = await create_test_file(server, "status.txt", source.id, default_user)
agent = await create_test_agent_with_files(server, "status-agent", default_user, [(source.id, file_metadata.id)])
exported = await agent_file_manager.export([agent.id], default_user)
exported = await agent_serialization_manager.export([agent.id], default_user)
result = await agent_file_manager.import_file(exported, other_user)
result = await agent_serialization_manager.import_file(exported, other_user)
imported_file_id = next(db_id for file_id, db_id in result.id_mappings.items() if file_id.startswith("file-"))
imported_file = await server.file_manager.get_file_by_id(imported_file_id, other_user)
@@ -1100,7 +1100,7 @@ class TestAgentFileImportWithProcessing:
assert imported_file.total_chunks is None
assert imported_file.chunks_embedded is None
async def test_import_multiple_files_processing(self, server, agent_file_manager, default_user, other_user):
async def test_import_multiple_files_processing(self, server, agent_serialization_manager, default_user, other_user):
"""Test import processes multiple files efficiently."""
source = await create_test_source(server, "multi-source", default_user)
file1 = await create_test_file(server, "file1.txt", source.id, default_user)
@@ -1108,9 +1108,9 @@ class TestAgentFileImportWithProcessing:
agent = await create_test_agent_with_files(server, "multi-agent", default_user, [(source.id, file1.id), (source.id, file2.id)])
exported = await agent_file_manager.export([agent.id], default_user)
exported = await agent_serialization_manager.export([agent.id], default_user)
result = await agent_file_manager.import_file(exported, other_user)
result = await agent_serialization_manager.import_file(exported, other_user)
imported_file_ids = [db_id for file_id, db_id in result.id_mappings.items() if file_id.startswith("file-")]
assert len(imported_file_ids) == 2
@@ -1123,24 +1123,24 @@ class TestAgentFileImportWithProcessing:
class TestAgentFileRoundTrip:
"""Tests for complete export -> import -> export cycles."""
async def test_roundtrip_consistency(self, server, agent_file_manager, test_agent, default_user, other_user):
async def test_roundtrip_consistency(self, server, agent_serialization_manager, test_agent, default_user, other_user):
"""Test that export -> import -> export produces consistent results."""
original_export = await agent_file_manager.export([test_agent.id], default_user)
result = await agent_file_manager.import_file(original_export, other_user)
original_export = await agent_serialization_manager.export([test_agent.id], default_user)
result = await agent_serialization_manager.import_file(original_export, other_user)
imported_agent_id = next(db_id for file_id, db_id in result.id_mappings.items() if file_id == "agent-0")
second_export = await agent_file_manager.export([imported_agent_id], other_user)
second_export = await agent_serialization_manager.export([imported_agent_id], other_user)
assert compare_agent_files(original_export, second_export)
async def test_multiple_roundtrips(self, server, agent_file_manager, test_agent, default_user, other_user):
async def test_multiple_roundtrips(self, server, agent_serialization_manager, test_agent, default_user, other_user):
"""Test multiple rounds of export/import maintain consistency."""
current_agent_id = test_agent.id
current_user = default_user
for i in range(3):
agent_file = await agent_file_manager.export([current_agent_id], current_user)
agent_file = await agent_serialization_manager.export([current_agent_id], current_user)
target_user = other_user if current_user == default_user else default_user
result = await agent_file_manager.import_file(agent_file, target_user)
result = await agent_serialization_manager.import_file(agent_file, target_user)
current_agent_id = next(db_id for file_id, db_id in result.id_mappings.items() if file_id == "agent-0")
current_user = target_user
@@ -1152,7 +1152,7 @@ class TestAgentFileRoundTrip:
class TestAgentFileEdgeCases:
"""Tests for edge cases and error conditions."""
async def test_agent_with_no_messages(self, server, agent_file_manager, default_user, other_user):
async def test_agent_with_no_messages(self, server, agent_serialization_manager, default_user, other_user):
"""Test exporting/importing agent with no messages."""
# Create agent with no initial messages
create_agent_request = CreateAgent(
@@ -1170,10 +1170,10 @@ class TestAgentFileEdgeCases:
)
# Export
agent_file = await agent_file_manager.export([agent_state.id], default_user)
agent_file = await agent_serialization_manager.export([agent_state.id], default_user)
# Import
result = await agent_file_manager.import_file(agent_file, other_user)
result = await agent_serialization_manager.import_file(agent_file, other_user)
# Verify
assert result.success
@@ -1182,7 +1182,7 @@ class TestAgentFileEdgeCases:
assert len(imported_agent.message_ids) == 0
async def test_large_agent_file(self, server, agent_file_manager, default_user, other_user, weather_tool):
async def test_large_agent_file(self, server, agent_serialization_manager, default_user, other_user, weather_tool):
"""Test handling of larger agent files with many messages."""
# Create agent
create_agent_request = CreateAgent(
@@ -1207,14 +1207,14 @@ class TestAgentFileEdgeCases:
)
# Export
agent_file = await agent_file_manager.export([agent_state.id], default_user)
agent_file = await agent_serialization_manager.export([agent_state.id], default_user)
# Verify large file
exported_agent = agent_file.agents[0]
assert len(exported_agent.messages) >= 10
# Import
result = await agent_file_manager.import_file(agent_file, other_user)
result = await agent_serialization_manager.import_file(agent_file, other_user)
# Verify all messages imported correctly
assert result.success