From c3a2e1f4be7f8ff8f053dd298656513281ac78d2 Mon Sep 17 00:00:00 2001 From: Kevin Lin Date: Thu, 23 Jan 2025 17:36:54 -0800 Subject: [PATCH] fix: fix the offline agent test (#761) --- letta/offline_memory_agent.py | 20 +++++++++---------- .../integration_test_offline_memory_agent.py | 8 ++++++-- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/letta/offline_memory_agent.py b/letta/offline_memory_agent.py index 076e2dc0..be8b9b33 100644 --- a/letta/offline_memory_agent.py +++ b/letta/offline_memory_agent.py @@ -8,12 +8,12 @@ from letta.schemas.openai.chat_completion_response import UsageStatistics from letta.schemas.usage import LettaUsageStatistics -def trigger_rethink_memory(agent_state: "AgentState", message: Optional[str]) -> Optional[str]: # type: ignore +def trigger_rethink_memory(agent_state: "AgentState", message: str) -> None: # type: ignore """ Called if and only when user says the word trigger_rethink_memory". It will trigger the re-evaluation of the memory. Args: - message (Optional[str]): Description of what aspect of the memory should be re-evaluated. + message (str): Description of what aspect of the memory should be re-evaluated. """ from letta import create_client @@ -25,12 +25,12 @@ def trigger_rethink_memory(agent_state: "AgentState", message: Optional[str]) -> client.user_message(agent_id=agent.id, message=message) -def trigger_rethink_memory_convo(agent_state: "AgentState", message: Optional[str]) -> Optional[str]: # type: ignore +def trigger_rethink_memory_convo(agent_state: "AgentState", message: str) -> None: # type: ignore """ Called if and only when user says the word "trigger_rethink_memory". It will trigger the re-evaluation of the memory. Args: - message (Optional[str]): Description of what aspect of the memory should be re-evaluated. + message (str): Description of what aspect of the memory should be re-evaluated. """ from letta import create_client @@ -48,7 +48,7 @@ def trigger_rethink_memory_convo(agent_state: "AgentState", message: Optional[st client.user_message(agent_id=agent.id, message=message) -def rethink_memory_convo(agent_state: "AgentState", new_memory: str, target_block_label: Optional[str], source_block_label: Optional[str]) -> Optional[str]: # type: ignore +def rethink_memory_convo(agent_state: "AgentState", new_memory: str, target_block_label: str, source_block_label: str) -> None: # type: ignore """ Re-evaluate the memory in block_name, integrating new and updated facts. Replace outdated information with the most likely truths, avoiding redundancy with original memories. Ensure consistency with other memory blocks. @@ -58,7 +58,7 @@ def rethink_memory_convo(agent_state: "AgentState", new_memory: str, target_bloc target_block_label (str): The name of the block to write to. This should be chat_agent_human_new or chat_agent_persona_new. Returns: - Optional[str]: None is always returned as this function does not produce a response. + None: None is always returned as this function does not produce a response. """ if target_block_label is not None: if agent_state.memory.get_block(target_block_label) is None: @@ -67,7 +67,7 @@ def rethink_memory_convo(agent_state: "AgentState", new_memory: str, target_bloc return None -def rethink_memory(agent_state: "AgentState", new_memory: str, target_block_label: Optional[str], source_block_label: Optional[str]) -> Optional[str]: # type: ignore +def rethink_memory(agent_state: "AgentState", new_memory: str, target_block_label: str, source_block_label: str) -> None: # type: ignore """ Re-evaluate the memory in block_name, integrating new and updated facts. Replace outdated information with the most likely truths, avoiding redundancy with original memories. @@ -78,7 +78,7 @@ def rethink_memory(agent_state: "AgentState", new_memory: str, target_block_labe source_block_label (str): The name of the block to integrate information from. None if all the information has been integrated to terminate the loop. target_block_label (str): The name of the block to write to. Returns: - Optional[str]: None is always returned as this function does not produce a response. + None: None is always returned as this function does not produce a response. """ if target_block_label is not None: @@ -88,7 +88,7 @@ def rethink_memory(agent_state: "AgentState", new_memory: str, target_block_labe return None -def finish_rethinking_memory(agent_state: "AgentState") -> Optional[str]: # type: ignore +def finish_rethinking_memory(agent_state: "AgentState") -> None: # type: ignore """ This function is called when the agent is done rethinking the memory. @@ -98,7 +98,7 @@ def finish_rethinking_memory(agent_state: "AgentState") -> Optional[str]: # typ return None -def finish_rethinking_memory_convo(agent_state: "AgentState") -> Optional[str]: # type: ignore +def finish_rethinking_memory_convo(agent_state: "AgentState") -> None: # type: ignore """ This function is called when the agent is done rethinking the memory. diff --git a/tests/integration_test_offline_memory_agent.py b/tests/integration_test_offline_memory_agent.py index 44ef23d9..e1214466 100644 --- a/tests/integration_test_offline_memory_agent.py +++ b/tests/integration_test_offline_memory_agent.py @@ -75,8 +75,10 @@ def test_ripple_edit(client, mock_e2b_api_key_none): # limit=2000, # ) # new_memory = Block(name="rethink_memory_block", label="rethink_memory_block", value="[empty]", limit=2000) - conversation_memory = BasicBlockMemory(blocks=[conversation_persona_block, conversation_human_block, fact_block, new_memory]) - offline_memory = BasicBlockMemory(blocks=[offline_persona_block, offline_human_block, fact_block, new_memory]) + # conversation_memory = BasicBlockMemory(blocks=[conversation_persona_block, conversation_human_block, fact_block, new_memory]) + conversation_memory = BasicBlockMemory(blocks=[conversation_persona_block, conversation_human_block, fact_block]) + # offline_memory = BasicBlockMemory(blocks=[offline_persona_block, offline_human_block, fact_block, new_memory]) + offline_memory = BasicBlockMemory(blocks=[offline_persona_block, offline_human_block, fact_block]) conversation_agent = client.create_agent( name="conversation_agent", @@ -86,6 +88,7 @@ def test_ripple_edit(client, mock_e2b_api_key_none): embedding_config=EmbeddingConfig.default_config("text-embedding-ada-002"), tool_ids=[send_message.id, trigger_rethink_memory_tool.id], memory=conversation_memory, + block_ids=[new_memory.id], include_base_tools=False, ) assert conversation_agent is not None @@ -103,6 +106,7 @@ def test_ripple_edit(client, mock_e2b_api_key_none): embedding_config=EmbeddingConfig.default_config("text-embedding-ada-002"), tool_ids=[rethink_memory_tool.id, finish_rethinking_memory_tool.id], tool_rules=[TerminalToolRule(tool_name=finish_rethinking_memory_tool.name)], + block_ids=[new_memory.id], include_base_tools=False, ) assert offline_memory_agent is not None