Spaces:
Running
Running
File size: 3,084 Bytes
b2d9e47 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import os
import json
import urllib.parse
import sys
# Append local path so we can import StateKV
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from db import StateKV
def import_old_data(old_db_path: str, kv: StateKV) -> bool:
if not os.path.exists(old_db_path):
print(f"[import] Error: Path {old_db_path} does not exist.")
return False
print(f"[import] Starting migration from: {old_db_path}")
count = 0
for filename in os.listdir(old_db_path):
if not filename.endswith(".bin"):
continue
if not filename.startswith("mem%3A"):
continue
# URL decode the scope name (e.g. mem%3Asessions.bin -> mem:sessions)
decoded_name = urllib.parse.unquote(filename)
scope = decoded_name[:-4] if decoded_name.endswith(".bin") else decoded_name
filepath = os.path.join(old_db_path, filename)
try:
with open(filepath, "rb") as f:
data = f.read()
if not data:
continue
# Find the last matching JSON curly brace to strip binary padding/footers
try:
last_brace = data.rindex(b'}')
json_str = data[:last_brace+1].decode("utf-8")
records = json.loads(json_str)
except Exception as e:
print(f"[import] Skipping {filename}: could not parse JSON ({e})")
continue
print(f"[import] Scope '{scope}': found {len(records)} items. Migrating...")
for key, val in records.items():
kv.set(scope, key, val)
count += 1
except Exception as e:
print(f"[import] Failed to process {filename}: {e}")
if count > 0:
print("[import] Creating Dolt version commit...")
kv.commit_version("Import legacy AgentMemory database from Hugging Face", "system")
print(f"[import] Finished! Migrated {count} total records into Dolt SQL.")
return True
if __name__ == "__main__":
# Prioritize the restored HF data in the user's home folder, falling back to local workspace data
home_path = os.path.expanduser(os.path.join("~", ".agentmemory", "state_store.db"))
workspace_path = r"d:\Downloads\Projects\Other Projects\agentmemory\agentmemory\data\state_store.db"
if os.path.exists(home_path):
default_path = home_path
print(f"[import] Found restored HF data at: {default_path}")
else:
default_path = workspace_path
print(f"[import] HF data not found at home path, falling back to workspace: {default_path}")
# Allow custom path via command line argument
if len(sys.argv) > 1:
default_path = sys.argv[1]
print(f"[import] Overriding path with CLI argument: {default_path}")
# Initialize connection
print(f"[import] Initializing Dolt StateKV...")
kv = StateKV()
import_old_data(default_path, kv)
|