From 198009c78468d0cdf68bae66ff9103edfe7dae07 Mon Sep 17 00:00:00 2001
From: Sarah Wooders <sarahwooders@gmail.com>
Date: Tue, 16 Jul 2024 22:38:30 -0700
Subject: [PATCH] fix: fix bug in migration script memory for 0.3.18 (#1548)

---
 examples/google_search_preset.yaml | 11 -----------
 scripts/migrate_0.3.18.py          | 10 ++++++++--
 2 files changed, 8 insertions(+), 13 deletions(-)
 delete mode 100644 examples/google_search_preset.yaml

diff --git a/examples/google_search_preset.yaml b/examples/google_search_preset.yaml
deleted file mode 100644
index 2454e1f2..00000000
--- a/examples/google_search_preset.yaml
+++ /dev/null
@@ -1,11 +0,0 @@
-system_prompt: "memgpt_chat"
-functions:
-  - "send_message"
-  - "pause_heartbeats"
-  - "core_memory_append"
-  - "core_memory_replace"
-  - "conversation_search"
-  - "conversation_search_date"
-  - "archival_memory_insert"
-  - "archival_memory_search"
-  - "google_search"
diff --git a/scripts/migrate_0.3.18.py b/scripts/migrate_0.3.18.py
index 5ac681ab..7bbd8d00 100644
--- a/scripts/migrate_0.3.18.py
+++ b/scripts/migrate_0.3.18.py
@@ -32,10 +32,10 @@ agent_states = []
 with engine.connect() as conn:
     agents = conn.execute(table.select()).fetchall()
     for agent in agents:
-        # print(agent)
         id = uuid.UUID(agent[0])
         user_id = uuid.UUID(agent[1])
         name = agent[2]
+        print(f"Migrating agent {name}")
         persona = agent[3]
         human = agent[4]
         system = agent[5]
@@ -46,7 +46,7 @@ with engine.connect() as conn:
         state = agent[10]
         tools = agent[11]
 
-        state["memory"] = {"human": human, "persona": persona}
+        state["memory"] = {"human": {"value": human, "limit": 2000}, "persona": {"value": persona, "limit": 2000}}
 
         agent_state = AgentState(
             id=id,
@@ -67,9 +67,15 @@ with engine.connect() as conn:
 agents_model = Table("agents", metadata, autoload_with=engine)
 agents_model.drop(engine)
 
+# remove tool table
+tool_model = Table("toolmodel", metadata, autoload_with=engine)
+tool_model.drop(engine)
+
 # re-create tables and add default tools
 ms = MetadataStore(config)
 add_default_tools(None, ms)
+print("Tools", [tool.name for tool in ms.list_tools()])
+
 
 for agent in agent_states:
     ms.create_agent(agent)