From 1e77c066189859646c79d9c42c4afa478aebd523 Mon Sep 17 00:00:00 2001
From: Sarah Wooders <sarahwooders@gmail.com>
Date: Thu, 11 Jul 2024 18:21:42 -0700
Subject: [PATCH] fix examples

---
 examples/memgpt_local_client.py               | 27 +++++++++++++++++++
 ...memgpt_client.py => memgpt_rest_client.py} | 26 ++++++------------
 memgpt/client/client.py                       |  2 +-
 3 files changed, 36 insertions(+), 19 deletions(-)
 create mode 100644 examples/memgpt_local_client.py
 rename examples/{memgpt_client.py => memgpt_rest_client.py} (55%)

diff --git a/examples/memgpt_local_client.py b/examples/memgpt_local_client.py
new file mode 100644
index 00000000..92624705
--- /dev/null
+++ b/examples/memgpt_local_client.py
@@ -0,0 +1,27 @@
+import json
+
+from memgpt import create_client
+from memgpt.memory import ChatMemory
+
+
+def main():
+
+    # Create a `LocalClient`
+    client = create_client()
+
+    # Create an agent
+    agent_state = client.create_agent(name="my_agent", memory=ChatMemory(human="My name is Sarah.", persona="I am a friendly AI."))
+    print(f"Created agent: {agent_state.name} with ID {str(agent_state.id)}")
+
+    # Send a message to the agent
+    print(f"Created agent: {agent_state.name} with ID {str(agent_state.id)}")
+    send_message_response = client.user_message(agent_id=agent_state.id, message="Whats my name?")
+    print(f"Recieved response: \n{json.dumps(send_message_response.messages, indent=4)}")
+
+    # Delete agent
+    client.delete_agent(agent_id=agent_state.id)
+    print(f"Deleted agent: {agent_state.name} with ID {str(agent_state.id)}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/memgpt_client.py b/examples/memgpt_rest_client.py
similarity index 55%
rename from examples/memgpt_client.py
rename to examples/memgpt_rest_client.py
index 06577060..9e524283 100644
--- a/examples/memgpt_client.py
+++ b/examples/memgpt_rest_client.py
@@ -1,8 +1,7 @@
 import json
 
 from memgpt import Admin, create_client
-from memgpt.constants import DEFAULT_HUMAN, DEFAULT_PERSONA, DEFAULT_PRESET
-from memgpt.utils import get_human_text, get_persona_text
+from memgpt.memory import ChatMemory
 
 """
 Make sure you run the MemGPT server before running this example.
@@ -31,26 +30,17 @@ def main():
     client = create_client(base_url="http://localhost:8283", token=token)
 
     # Create an agent
-    agent_info = client.create_agent(
-        name="my_agent",
-        preset=DEFAULT_PRESET,
-        persona=get_persona_text(DEFAULT_PERSONA),
-        human=get_human_text(DEFAULT_HUMAN),
-    )
-    print(f"Created agent: {agent_info.name} with ID {str(agent_info.id)}")
+    agent_state = client.create_agent(name="my_agent", memory=ChatMemory(human="My name is Sarah.", persona="I am a friendly AI."))
+    print(f"Created agent: {agent_state.name} with ID {str(agent_state.id)}")
 
     # Send a message to the agent
-    send_message_response = client.user_message(agent_id=agent_info.id, message="Hello, agent!")
-    messages = send_message_response.messages
-    print(f"Recieved response: \n{json.dumps(messages, indent=4)}")
-
-    # TODO: get agent memory
-
-    # TODO: Update agent persona
+    print(f"Created agent: {agent_state.name} with ID {str(agent_state.id)}")
+    send_message_response = client.user_message(agent_id=agent_state.id, message="Whats my name?")
+    print(f"Recieved response: \n{json.dumps(send_message_response.messages, indent=4)}")
 
     # Delete agent
-    client.delete_agent(agent_id=agent_info.id)
-    print(f"Deleted agent: {agent_info.name} with ID {str(agent_info.id)}")
+    client.delete_agent(agent_id=agent_state.id)
+    print(f"Deleted agent: {agent_state.name} with ID {str(agent_state.id)}")
 
     # Delete user
     admin.delete_user(user_id=user_id)
diff --git a/memgpt/client/client.py b/memgpt/client/client.py
index 2b2bc316..97393207 100644
--- a/memgpt/client/client.py
+++ b/memgpt/client/client.py
@@ -811,7 +811,7 @@ class LocalClient(AbstractClient):
         else:
             return UserMessageResponse(messages=self.interface.to_list(), usage=usage)
 
-    def user_message(self, agent_id: str, message: str) -> Union[List[Dict], Tuple[List[Dict], int]]:
+    def user_message(self, agent_id: str, message: str) -> UserMessageResponse:
         self.interface.clear()
         usage = self.server.user_message(user_id=self.user_id, agent_id=agent_id, message=message)
         if self.auto_save: