From 8d7bd0828e3ba55c2ce302b6a63201d3a3b64268 Mon Sep 17 00:00:00 2001
From: Sarah Wooders <sarahwooders@gmail.com>
Date: Tue, 30 Jul 2024 22:23:06 -0700
Subject: [PATCH] feat: various fixes to improve notebook useability  (#1593)

---
 memgpt/client/client.py | 19 +++++++++++++++++--
 memgpt/server/server.py |  2 +-
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/memgpt/client/client.py b/memgpt/client/client.py
index e92c0e9b1c..cfe0731379 100644
--- a/memgpt/client/client.py
+++ b/memgpt/client/client.py
@@ -771,10 +771,13 @@ def rename_agent(self, agent_id: uuid.UUID, new_name: str):
     def delete_agent(self, agent_id: uuid.UUID):
         self.server.delete_agent(user_id=self.user_id, agent_id=agent_id)
 
-    def get_agent(self, agent_id: uuid.UUID) -> AgentState:
+    def get_agent_config(self, agent_id: uuid.UUID) -> AgentState:
         self.interface.clear()
         return self.server.get_agent_config(user_id=self.user_id, agent_id=agent_id)
 
+    def get_agent(self, agent_id: Optional[uuid.UUID] = None, agent_name: Optional[str] = None):
+        return self.server.ms.get_agent(user_id=self.user_id, agent_id=agent_id, agent_name=agent_name)
+
     # presets
     def create_preset(self, preset: Preset) -> Preset:
         if preset.user_id is None:
@@ -799,7 +802,19 @@ def update_agent_core_memory(self, agent_id: str, new_memory_contents: Dict) ->
 
     # agent interactions
 
-    def send_message(self, agent_id: uuid.UUID, message: str, role: str, stream: Optional[bool] = False) -> UserMessageResponse:
+    def send_message(
+        self,
+        message: str,
+        role: str,
+        agent_id: Optional[uuid.UUID] = None,
+        agent_name: Optional[str] = None,
+        stream: Optional[bool] = False,
+    ) -> UserMessageResponse:
+        if not agent_id:
+            assert agent_name, f"Either agent_id or agent_name must be provided"
+            agent_state = self.get_agent(agent_name=agent_name)
+            agent_id = agent_state.id
+
         if stream:
             # TODO: implement streaming with stream=True/False
             raise NotImplementedError
diff --git a/memgpt/server/server.py b/memgpt/server/server.py
index 9d39ff1897..377546df76 100644
--- a/memgpt/server/server.py
+++ b/memgpt/server/server.py
@@ -653,7 +653,7 @@ def system_message(
             message.created_at = timestamp
 
         # Run the agent state forward
-        return self._step(user_id=user_id, agent_id=agent_id, input_message=packaged_system_message, timestamp=None)
+        return self._step(user_id=user_id, agent_id=agent_id, input_message=packaged_system_message, timestamp=timestamp)
 
     # @LockingServer.agent_lock_decorator
     def run_command(self, user_id: uuid.UUID, agent_id: uuid.UUID, command: str) -> Union[MemGPTUsageStatistics, None]: