diff --git a/letta/services/agent_manager.py b/letta/services/agent_manager.py index fd3afcfd..d2f4e88b 100644 --- a/letta/services/agent_manager.py +++ b/letta/services/agent_manager.py @@ -38,7 +38,7 @@ from letta.schemas.group import ManagerType from letta.schemas.llm_config import LLMConfig from letta.schemas.memory import Memory from letta.schemas.message import Message as PydanticMessage -from letta.schemas.message import MessageCreate +from letta.schemas.message import MessageCreate, MessageUpdate from letta.schemas.passage import Passage as PydanticPassage from letta.schemas.source import Source as PydanticSource from letta.schemas.tool import Tool as PydanticTool @@ -718,10 +718,12 @@ class AgentManager: model=agent_state.llm_config.model, openai_message_dict={"role": "system", "content": new_system_message_str}, ) - # TODO: This seems kind of silly, why not just update the message? - message = self.message_manager.create_message(message, actor=actor) - message_ids = [message.id] + agent_state.message_ids[1:] # swap index 0 (system) - return self.set_in_context_messages(agent_id=agent_id, message_ids=message_ids, actor=actor) + message = self.message_manager.update_message_by_id( + message_id=curr_system_message.id, + message_update=MessageUpdate(**message.model_dump()), + actor=actor, + ) + return self.set_in_context_messages(agent_id=agent_id, message_ids=agent_state.message_ids, actor=actor) else: return agent_state diff --git a/tests/test_server.py b/tests/test_server.py index 51bd5816..dd128913 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -279,8 +279,6 @@ def server(): def org_id(server): # create org org = server.organization_manager.create_default_organization() - print(f"Created org\n{org.id}") - yield org.id # cleanup @@ -338,7 +336,6 @@ def agent_id(server, user_id, base_tools): ), actor=actor, ) - print(f"Created agent\n{agent_state}") yield agent_state.id # cleanup @@ -359,7 +356,6 @@ def other_agent_id(server, user_id, base_tools): ), actor=actor, ) - print(f"Created agent\n{agent_state}") yield agent_state.id # cleanup @@ -953,7 +949,6 @@ def test_memory_rebuild_count(server, user, disable_e2b_api_key, base_tools, bas ), actor=actor, ) - print(f"Created agent\n{agent_state}") def count_system_messages_in_recall() -> Tuple[int, List[LettaMessage]]: @@ -967,10 +962,6 @@ def test_memory_rebuild_count(server, user, disable_e2b_api_key, base_tools, bas ) assert all(isinstance(m, LettaMessage) for m in letta_messages) - print("LETTA_MESSAGES:") - for i, m in enumerate(letta_messages): - print(f"{i}: {type(m)} ...{str(m)[-50:]}") - # Collect system messages and their texts system_messages = [m for m in letta_messages if m.message_type == "system_message"] return len(system_messages), letta_messages