From eb26f401b6c8ee8a603b0e5c822c95088fa5ec4b Mon Sep 17 00:00:00 2001 From: Kevin Lin Date: Fri, 6 Dec 2024 11:12:35 -0800 Subject: [PATCH 1/5] check if block is in labels --- letta/offline_memory_agent.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/letta/offline_memory_agent.py b/letta/offline_memory_agent.py index 16447c4614..c4b6d8487f 100644 --- a/letta/offline_memory_agent.py +++ b/letta/offline_memory_agent.py @@ -61,7 +61,7 @@ def rethink_memory_convo(agent_state: "AgentState", new_memory: str, target_bloc Optional[str]: 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: + if target_block_label not in agent_state.memory.list_block_labels(): 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) return None @@ -82,7 +82,7 @@ def rethink_memory(agent_state: "AgentState", new_memory: str, target_block_labe """ if target_block_label is not None: - if agent_state.memory.get_block(target_block_label) is None: + if target_block_label not in agent_state.memory.list_block_labels(): 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) return None From ec888ea461f2f4391124debdeb2d7c57484e081e Mon Sep 17 00:00:00 2001 From: Kevin Lin Date: Fri, 6 Dec 2024 12:04:10 -0800 Subject: [PATCH 2/5] fix block creation --- letta/offline_memory_agent.py | 18 +++++++-- tests/test_offline_memory_agent.py | 59 +++++++++++++++++++++++------- 2 files changed, 59 insertions(+), 18 deletions(-) diff --git a/letta/offline_memory_agent.py b/letta/offline_memory_agent.py index c4b6d8487f..6011ba2844 100644 --- a/letta/offline_memory_agent.py +++ b/letta/offline_memory_agent.py @@ -18,6 +18,7 @@ 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: @@ -60,10 +61,15 @@ 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 target_block_label not in agent_state.memory.list_block_labels(): - 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) + 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 @@ -80,11 +86,15 @@ 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 target_block_label not in agent_state.memory.list_block_labels(): - 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) + 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 diff --git a/tests/test_offline_memory_agent.py b/tests/test_offline_memory_agent.py index b78e2274c8..31ff1f2865 100644 --- a/tests/test_offline_memory_agent.py +++ b/tests/test_offline_memory_agent.py @@ -1,26 +1,55 @@ -import json +import pytest from letta import BasicBlockMemory -from letta import offline_memory_agent from letta.client.client import Block, create_client from letta.constants import DEFAULT_HUMAN, DEFAULT_PERSONA from letta.offline_memory_agent import ( - rethink_memory, finish_rethinking_memory, - rethink_memory_convo, finish_rethinking_memory_convo, + rethink_memory, + rethink_memory_convo, trigger_rethink_memory, - trigger_rethink_memory_convo, ) from letta.prompts import gpt_system from letta.schemas.agent import AgentType from letta.schemas.embedding_config import EmbeddingConfig from letta.schemas.llm_config import LLMConfig -from letta.schemas.message import MessageCreate from letta.schemas.tool_rule import TerminalToolRule from letta.utils import get_human_text, get_persona_text +@pytest.fixture(scope="module") +def client(): + client = create_client() + client = create_client() + assert client is not None + client.set_default_llm_config(LLMConfig.default_config("gpt-4o-mini")) + client.create_tool(trigger_rethink_memory) + client.set_default_embedding_config(EmbeddingConfig.default_config(provider="openai")) + + yield client + + +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 = create_client() assert client is not None @@ -61,12 +90,14 @@ def test_ripple_edit(): ) assert conversation_agent is not None - assert set(conversation_agent.memory.list_block_labels()) == set([ - "persona", - "human", - "fact_block", - "rethink_memory_block", - ]) + assert set(conversation_agent.memory.list_block_labels()) == set( + [ + "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) @@ -82,7 +113,7 @@ def test_ripple_edit(): include_base_tools=False, ) assert offline_memory_agent is not None - assert set(offline_memory_agent.memory.list_block_labels())== set(["persona", "human", "fact_block", "rethink_memory_block"]) + assert set(offline_memory_agent.memory.list_block_labels()) == set(["persona", "human", "fact_block", "rethink_memory_block"]) response = client.user_message( agent_id=conversation_agent.id, message="[trigger_rethink_memory]: Messi has now moved to playing for Inter Miami" ) @@ -114,7 +145,7 @@ def test_chat_only_agent(): tools=["send_message"], memory=conversation_memory, include_base_tools=False, - metadata = {"offline_memory_tools": [rethink_memory.name, finish_rethinking_memory.name]} + metadata={"offline_memory_tools": [rethink_memory.name, finish_rethinking_memory.name]}, ) assert chat_only_agent is not None assert set(chat_only_agent.memory.list_block_labels()) == set(["chat_agent_persona", "chat_agent_human"]) From 9148909d2d9c3fbe9529e5afd1004d2dc02e1b0f Mon Sep 17 00:00:00 2001 From: Kevin Lin Date: Fri, 6 Dec 2024 13:26:38 -0800 Subject: [PATCH 3/5] clean up --- tests/test_offline_memory_agent.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/test_offline_memory_agent.py b/tests/test_offline_memory_agent.py index 740354203b..de6c91da53 100644 --- a/tests/test_offline_memory_agent.py +++ b/tests/test_offline_memory_agent.py @@ -8,8 +8,6 @@ finish_rethinking_memory_convo, rethink_memory, rethink_memory_convo, - rethink_memory, - rethink_memory_convo, trigger_rethink_memory, ) from letta.prompts import gpt_system @@ -22,7 +20,6 @@ @pytest.fixture(scope="module") def client(): - client = create_client() client = create_client() assert client is not None client.set_default_llm_config(LLMConfig.default_config("gpt-4o-mini")) @@ -152,7 +149,6 @@ def test_chat_only_agent(client, mock_e2b_api_key_none): memory=conversation_memory, include_base_tools=False, metadata={"offline_memory_tools": [rethink_memory.name, finish_rethinking_memory.name]}, - metadata={"offline_memory_tools": [rethink_memory.name, finish_rethinking_memory.name]}, ) assert chat_only_agent is not None assert set(chat_only_agent.memory.list_block_labels()) == {"chat_agent_persona", "chat_agent_human"} From 2d6eed8046b72d0c98ea1e263cc7fe167733a99b Mon Sep 17 00:00:00 2001 From: Kevin Lin Date: Fri, 6 Dec 2024 13:43:00 -0800 Subject: [PATCH 4/5] clean up --- tests/test_offline_memory_agent.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/test_offline_memory_agent.py b/tests/test_offline_memory_agent.py index de6c91da53..5e3a2e6589 100644 --- a/tests/test_offline_memory_agent.py +++ b/tests/test_offline_memory_agent.py @@ -21,9 +21,7 @@ @pytest.fixture(scope="module") def client(): client = create_client() - assert client is not None client.set_default_llm_config(LLMConfig.default_config("gpt-4o-mini")) - client.create_tool(trigger_rethink_memory) client.set_default_embedding_config(EmbeddingConfig.default_config(provider="openai")) yield client @@ -91,14 +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()) == set( - [ + 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) From 617b24f5f176ed30a6c76d4d8fd4f547840b16a8 Mon Sep 17 00:00:00 2001 From: Kevin Lin Date: Fri, 6 Dec 2024 14:01:33 -0800 Subject: [PATCH 5/5] clean up --- letta/offline_memory_agent.py | 5 +---- tests/test_offline_memory_agent.py | 2 +- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/letta/offline_memory_agent.py b/letta/offline_memory_agent.py index e23030d8ac..6cf90ab092 100644 --- a/letta/offline_memory_agent.py +++ b/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: @@ -62,7 +60,6 @@ def rethink_memory_convo(agent_state: "AgentState", new_memory: str, target_bloc 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 target_block_label not in agent_state.memory.list_block_labels(): @@ -87,7 +84,6 @@ def rethink_memory(agent_state: "AgentState", new_memory: str, target_block_labe 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 target_block_label not in agent_state.memory.list_block_labels(): @@ -98,6 +94,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 """ This function is called when the agent is done rethinking the memory. diff --git a/tests/test_offline_memory_agent.py b/tests/test_offline_memory_agent.py index 5e3a2e6589..225e76c98c 100644 --- a/tests/test_offline_memory_agent.py +++ b/tests/test_offline_memory_agent.py @@ -110,7 +110,7 @@ def test_ripple_edit(client, mock_e2b_api_key_none): include_base_tools=False, ) assert offline_memory_agent is not None - assert set(offline_memory_agent.memory.list_block_labels()) == set(["persona", "human", "fact_block", "rethink_memory_block"]) + assert set(offline_memory_agent.memory.list_block_labels()) == {"persona", "human", "fact_block", "rethink_memory_block"} response = client.user_message( agent_id=conversation_agent.id, message="[trigger_rethink_memory]: Messi has now moved to playing for Inter Miami" )