fix: fix rethink memory non existent block error

letta/offline_memory_agent.py

@@ -18,7 +18,6 @@ def trigger_rethink_memory(agent_state: "AgentState", message: Optional[str]) ->
 
     """
     from letta import create_client
-
     client = create_client()
     agents = client.list_agents()
     for agent in agents:
@@ -35,7 +34,6 @@ def trigger_rethink_memory_convo(agent_state: "AgentState", message: Optional[st
 
     """
     from letta import create_client
-
     client = create_client()
     recent_convo = "".join([str(message) for message in agent_state.messages])[
         -2000:
@@ -61,10 +59,14 @@ def rethink_memory_convo(agent_state: "AgentState", new_memory: str, target_bloc
     Returns:
         Optional[str]: None is always returned as this function does not produce a response.
     """
+    from letta import create_client
+    client = create_client()
     if target_block_label is not None:
-        if agent_state.memory.get_block(target_block_label) is None:
-            agent_state.memory.create_block(label=target_block_label, value=new_memory)
-        agent_state.memory.update_block_value(label=target_block_label, value=new_memory)
+        if target_block_label not in agent_state.memory.list_block_labels():
+            new_block = client.create_block(label=target_block_label, value=new_memory)
+            agent_state.memory.set_block(new_block)
+        else:
+            agent_state.memory.update_block_value(label=target_block_label, value=new_memory)
     return None
 
 
@@ -81,14 +83,18 @@ def rethink_memory(agent_state: "AgentState", new_memory: str, target_block_labe
     Returns:
         Optional[str]: None is always returned as this function does not produce a response.
     """
-
+    from letta import create_client
+    client = create_client()
     if target_block_label is not None:
-        if agent_state.memory.get_block(target_block_label) is None:
-            agent_state.memory.create_block(label=target_block_label, value=new_memory)
-        agent_state.memory.update_block_value(label=target_block_label, value=new_memory)
+        if target_block_label not in agent_state.memory.list_block_labels():
+            new_block = client.create_block(label=target_block_label, value=new_memory)
+            agent_state.memory.set_block(new_block)
+        else:
+            agent_state.memory.update_block_value(label=target_block_label, value=new_memory)
     return None
 
 
+
 def finish_rethinking_memory(agent_state: "AgentState") -> Optional[str]:  # type: ignore
     """
     This function is called when the agent is done rethinking the memory.

tests/test_offline_memory_agent.py

@@ -26,12 +26,30 @@ def client():
 
     yield client
 
-
 @pytest.fixture(autouse=True)
 def clear_agents(client):
     for agent in client.list_agents():
         client.delete_agent(agent.id)
 
+def test_rethink_memory_new_block(client):
+    """
+    Test that when rethink memory is called with a block that does not exist in the agent,
+    the new block is created.
+    """
+    client.create_agent()
+    agent = client.create_agent(
+        agent_type=AgentType.memgpt_agent,
+        system=gpt_system.get_system_text("memgpt_convo_only"),
+        llm_config=LLMConfig.default_config("gpt-4"),
+        embedding_config=EmbeddingConfig.default_config("text-embedding-ada-002"),
+        include_base_tools=False,
+    )
+    assert set(agent.memory.list_block_labels()) == {"persona", "human"}
+    rethink_memory(
+        agent_state=agent, new_memory="I am a new memory block content", source_block_label="human", target_block_label="new_memory_block"
+    )
+    assert set(agent.memory.list_block_labels()) == {"persona", "human", "new_memory_block"}
+
 
 def test_ripple_edit(client, mock_e2b_api_key_none):
     trigger_rethink_memory_tool = client.create_or_update_tool(trigger_rethink_memory)
@@ -71,7 +89,12 @@ def test_ripple_edit(client, mock_e2b_api_key_none):
     )
     assert conversation_agent is not None
 
-    assert set(conversation_agent.memory.list_block_labels()) == {"persona", "human", "fact_block", "rethink_memory_block"}
+    assert set(conversation_agent.memory.list_block_labels()) == { 
+            "persona",
+            "human",
+            "fact_block",
+            "rethink_memory_block",
+    }
 
     rethink_memory_tool = client.create_tool(rethink_memory)
     finish_rethinking_memory_tool = client.create_tool(finish_rethinking_memory)