From cee3c7ed580765940cd5957e749f05e8a060d1eb Mon Sep 17 00:00:00 2001 From: Ari Webb Date: Thu, 20 Nov 2025 13:53:42 -0800 Subject: [PATCH] fix: don't embed system messages [LET-6243] (#6307) don't embed system messages Co-authored-by: Ari Webb --- letta/services/message_manager.py | 18 +++++---- tests/integration_test_turbopuffer.py | 56 +++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 7 deletions(-) diff --git a/letta/services/message_manager.py b/letta/services/message_manager.py index 3da3f249..6cb4f5e9 100644 --- a/letta/services/message_manager.py +++ b/letta/services/message_manager.py @@ -537,13 +537,17 @@ class MessageManager: if should_use_tpuf_for_messages() and result: agent_id = result[0].agent_id if agent_id: - if strict_mode: - await self._embed_messages_background(result, actor, agent_id, project_id, template_id) - else: - fire_and_forget( - self._embed_messages_background(result, actor, agent_id, project_id, template_id), - task_name=f"embed_messages_for_agent_{agent_id}", - ) + # Filter out system messages before embedding to avoid unnecessary processing + # System messages (especially initial agent system messages) can be very large + messages_to_embed = [msg for msg in result if msg.role != MessageRole.system] + if messages_to_embed: + if strict_mode: + await self._embed_messages_background(messages_to_embed, actor, agent_id, project_id, template_id) + else: + fire_and_forget( + self._embed_messages_background(messages_to_embed, actor, agent_id, project_id, template_id), + task_name=f"embed_messages_for_agent_{agent_id}", + ) if allow_partial and existing_messages: async with db_registry.async_session() as session: diff --git a/tests/integration_test_turbopuffer.py b/tests/integration_test_turbopuffer.py index 6387d95a..859f91f8 100644 --- a/tests/integration_test_turbopuffer.py +++ b/tests/integration_test_turbopuffer.py @@ -2165,3 +2165,59 @@ async def test_message_template_id_filtering(server, sarah_agent, default_user, await tpuf_client.delete_messages( agent_id=sarah_agent.id, organization_id=default_user.organization_id, message_ids=[message_a.id, message_b.id] ) + + +@pytest.mark.asyncio +async def test_system_messages_not_embedded_during_agent_creation(server, default_user, enable_message_embedding): + """Test that system messages are filtered out before being passed to the embedding pipeline during agent creation""" + from unittest.mock import AsyncMock, patch + + from letta.schemas.agent import CreateAgent + from letta.schemas.llm_config import LLMConfig + + # Mock the _embed_messages_background method to track what messages are passed to it + messages_passed_to_embed = [] + + original_embed = server.message_manager._embed_messages_background + + async def mock_embed(messages, actor, agent_id, project_id=None, template_id=None): + # Capture what messages are being passed to embedding + messages_passed_to_embed.extend(messages) + # Call the original method + return await original_embed(messages, actor, agent_id, project_id, template_id) + + with patch.object(server.message_manager, "_embed_messages_background", mock_embed): + # Create agent with initial messages (which includes a system message) + agent = await server.agent_manager.create_agent_async( + agent_create=CreateAgent( + name="TestSystemMessageAgent", + memory_blocks=[], + llm_config=LLMConfig.default_config("gpt-4o-mini"), + embedding_config=EmbeddingConfig.default_config(provider="openai"), + include_base_tools=False, + ), + actor=default_user, + ) + + # Get all messages created for the agent + all_messages = await server.message_manager.get_messages_by_ids_async(message_ids=agent.message_ids, actor=default_user) + + # Verify that at least one system message was created + system_messages = [msg for msg in all_messages if msg.role == MessageRole.system] + assert len(system_messages) > 0, "No system messages were created during agent creation" + + print(messages_passed_to_embed) + print(system_messages) + print(all_messages) + + # Verify that NO system messages were passed to the embedding pipeline + system_messages_in_embed = [msg for msg in messages_passed_to_embed if msg.role == MessageRole.system] + assert len(system_messages_in_embed) == 0, ( + f"System messages should not be embedded, but {len(system_messages_in_embed)} were passed to embedding pipeline" + ) + + # Clean up + try: + await server.agent_manager.delete_agent_async(agent.id, default_user) + except: + pass