Spaces:
Running
Running
| 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) | |