mirror of
https://github.com/cpacker/MemGPT.git
synced 2025-06-03 04:30:22 +00:00
85faf5f474
Co-authored-by: Charles Packer <packercharles@gmail.com> Co-authored-by: Shubham Naik <shubham.naik10@gmail.com> Co-authored-by: Shubham Naik <shub@memgpt.ai>
68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
import os
|
|
|
|
from sqlalchemy import DDL, MetaData, Table, create_engine, update
|
|
|
|
from letta.config import LettaConfig
|
|
from letta.constants import BASE_TOOLS
|
|
from letta.metadata import MetadataStore
|
|
from letta.presets.presets import add_default_tools
|
|
from letta.prompts import gpt_system
|
|
|
|
# Replace this with your actual database connection URL
|
|
config = LettaConfig.load()
|
|
if config.recall_storage_type == "sqlite":
|
|
DATABASE_URL = "sqlite:///" + os.path.join(config.recall_storage_path, "sqlite.db")
|
|
else:
|
|
DATABASE_URL = config.recall_storage_uri
|
|
print(DATABASE_URL)
|
|
engine = create_engine(DATABASE_URL)
|
|
metadata = MetaData()
|
|
|
|
# defaults
|
|
system_prompt = gpt_system.get_system_text("memgpt_chat")
|
|
|
|
# Reflect the existing table
|
|
table = Table("agents", metadata, autoload_with=engine)
|
|
|
|
# Using a connection to manage adding columns and committing updates
|
|
with engine.connect() as conn:
|
|
trans = conn.begin()
|
|
try:
|
|
# Check and add 'system' column if it does not exist
|
|
if "system" not in table.c:
|
|
ddl_system = DDL("ALTER TABLE agents ADD COLUMN system VARCHAR")
|
|
conn.execute(ddl_system)
|
|
# Reflect the table again to update metadata
|
|
metadata.clear()
|
|
table = Table("agents", metadata, autoload_with=conn)
|
|
|
|
# Check and add 'tools' column if it does not exist
|
|
if "tools" not in table.c:
|
|
ddl_tools = DDL("ALTER TABLE agents ADD COLUMN tools JSON")
|
|
conn.execute(ddl_tools)
|
|
# Reflect the table again to update metadata
|
|
metadata.clear()
|
|
table = Table("agents", metadata, autoload_with=conn)
|
|
|
|
# Update all existing rows with default values for the new columns
|
|
conn.execute(update(table).values(system=system_prompt, tools=BASE_TOOLS))
|
|
|
|
# Commit transaction
|
|
trans.commit()
|
|
print("Columns added and data updated successfully!")
|
|
|
|
except Exception as e:
|
|
print("An error occurred:", e)
|
|
trans.rollback() # Rollback if there are errors
|
|
|
|
# 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()])
|
|
|
|
print("Migration completed successfully!")
|