From 6094726dfe2de8906a6d8abf42b99374faabdfb2 Mon Sep 17 00:00:00 2001
From: Charles Packer <packercharles@gmail.com>
Date: Sat, 20 Jan 2024 17:28:39 -0800
Subject: [PATCH] fix: various fixes + make sure to save agent on creation to
 avoid missing agent_state.state entries (#877)

---
 memgpt/cli/cli.py       | 5 +++--
 memgpt/metadata.py      | 2 +-
 memgpt/server/server.py | 1 +
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/memgpt/cli/cli.py b/memgpt/cli/cli.py
index 8b4a395d..a5f082a6 100644
--- a/memgpt/cli/cli.py
+++ b/memgpt/cli/cli.py
@@ -28,7 +28,7 @@ from memgpt.agent import Agent
 from memgpt.embeddings import embedding_model
 from memgpt.server.constants import WS_DEFAULT_PORT, REST_DEFAULT_PORT
 from memgpt.data_types import AgentState, LLMConfig, EmbeddingConfig, User
-from memgpt.metadata import MetadataStore
+from memgpt.metadata import MetadataStore, save_agent
 from memgpt.migrate import migrate_all_agents, migrate_all_sources
 
 
@@ -604,6 +604,7 @@ def run(
                 agent_state=agent_state,
                 interface=interface,
             )
+            save_agent(agent=memgpt_agent, ms=ms)
         except ValueError as e:
             # TODO(swooders) what's the equivalent cleanup code for the new DB refactor?
             typer.secho(f"Failed to create agent from provided information:\n{e}", fg=typer.colors.RED)
@@ -623,7 +624,7 @@ def run(
             # except:
             #     typer.secho(f"Failed to delete agent directory during cleanup:\n{e}", fg=typer.colors.RED)
             sys.exit(1)
-        typer.secho(f"🎉 Created new agent '{agent_state.name}'", fg=typer.colors.GREEN)
+        typer.secho(f"🎉 Created new agent '{agent_state.name}' (id={agent_state.id})", fg=typer.colors.GREEN)
 
     # pretty print agent config
     # printd(json.dumps(vars(agent_config), indent=4, sort_keys=True, ensure_ascii=JSON_ENSURE_ASCII))
diff --git a/memgpt/metadata.py b/memgpt/metadata.py
index c1902543..3db2c757 100644
--- a/memgpt/metadata.py
+++ b/memgpt/metadata.py
@@ -64,7 +64,7 @@ class LLMConfigColumn(TypeDecorator):
         return value
 
     def process_result_value(self, value, dialect):
-        print("GET VALUE", value)
+        # print("GET VALUE", value)
         if value:
             return LLMConfig(**value)
         return value
diff --git a/memgpt/server/server.py b/memgpt/server/server.py
index 6c135088..ccee3be5 100644
--- a/memgpt/server/server.py
+++ b/memgpt/server/server.py
@@ -611,6 +611,7 @@ class SyncServer(LockingServer):
         logger.debug(f"Attempting to create agent from agent_state:\n{agent_state}")
         try:
             agent = presets.create_agent_from_preset(agent_state=agent_state, interface=interface)
+            save_agent(agent=agent, ms=self.ms)
 
             # FIXME: this is a hacky way to get the system prompts injected into agent into the DB
             # self.ms.update_agent(agent.agent_state)