MCP Memory Service

#!/usr/bin/env python3 # Copyright 2024 Heinrich Krupp # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ Enhanced migration script for MCP Memory Service. This script handles migration of memories between different ChromaDB instances, with support for both local and remote migrations. """ import sys import os from dotenv import load_dotenv from pathlib import Path import chromadb from chromadb import HttpClient, Settings import json import time from chromadb.utils import embedding_functions # Import our environment verifier from verify_environment import EnvironmentVerifier def verify_environment(): """Verify the environment before proceeding with migration""" verifier = EnvironmentVerifier() verifier.run_verifications() if not verifier.print_results(): print("\n⚠️ Environment verification failed! Migration cannot proceed.") sys.exit(1) print("\n✓ Environment verification passed! Proceeding with migration.") # Load environment variables load_dotenv() def get_claude_desktop_chroma_path(): """Get ChromaDB path from Claude Desktop config""" base_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) config_path = os.path.join(base_path, 'claude_config', 'mcp-memory', 'chroma_db') print(f"Using ChromaDB path: {config_path}") return config_path def migrate_memories(source_type, source_config, target_type, target_config): """ Migrate memories between ChromaDB instances. Args: source_type: 'local' or 'remote' source_config: For local: path to ChromaDB, for remote: {'host': host, 'port': port} target_type: 'local' or 'remote' target_config: For local: path to ChromaDB, for remote: {'host': host, 'port': port} """ print(f"Starting migration from {source_type} to {target_type}") try: # Set up embedding function embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction( model_name='all-MiniLM-L6-v2' ) # Connect to target ChromaDB if target_type == 'remote': target_client = HttpClient( host=target_config['host'], port=target_config['port'] ) print(f"Connected to remote ChromaDB at {target_config['host']}:{target_config['port']}") else: settings = Settings( anonymized_telemetry=False, allow_reset=True, is_persistent=True, persist_directory=target_config ) target_client = chromadb.Client(settings) print(f"Connected to local ChromaDB at {target_config}") # Get or create collection for imported memories try: target_collection = target_client.get_collection( name="mcp_imported_memories", embedding_function=embedding_function ) print("Found existing collection 'mcp_imported_memories' on target") except Exception: target_collection = target_client.create_collection( name="mcp_imported_memories", metadata={"hnsw:space": "cosine"}, embedding_function=embedding_function ) print("Created new collection 'mcp_imported_memories' on target") # Connect to source ChromaDB if source_type == 'remote': source_client = HttpClient( host=source_config['host'], port=source_config['port'] ) print(f"Connected to remote ChromaDB at {source_config['host']}:{source_config['port']}") else: settings = Settings( anonymized_telemetry=False, allow_reset=True, is_persistent=True, persist_directory=source_config ) source_client = chromadb.Client(settings) print(f"Connected to local ChromaDB at {source_config}") # List collections collections = source_client.list_collections() print(f"Found {len(collections)} collections in source") for coll in collections: print(f"- {coll.name}") # Try to get the memory collection try: source_collection = source_client.get_collection( name="memory_collection", embedding_function=embedding_function ) print("Found source memory collection") except ValueError as e: print(f"Error accessing source collection: {str(e)}") return # Get all memories from source print("Fetching source memories...") results = source_collection.get() if not results["ids"]: print("No memories found in source collection") return print(f"Found {len(results['ids'])} memories to migrate") # Check for existing memories in target to avoid duplicates target_existing = target_collection.get() existing_ids = set(target_existing["ids"]) # Filter out already migrated memories new_memories = { "ids": [], "documents": [], "metadatas": [] } for i, memory_id in enumerate(results["ids"]): if memory_id not in existing_ids: new_memories["ids"].append(memory_id) new_memories["documents"].append(results["documents"][i]) new_memories["metadatas"].append(results["metadatas"][i]) if not new_memories["ids"]: print("All memories are already migrated!") return print(f"Found {len(new_memories['ids'])} new memories to migrate") # Import in batches of 10 batch_size = 10 for i in range(0, len(new_memories['ids']), batch_size): batch_end = min(i + batch_size, len(new_memories['ids'])) batch_ids = new_memories['ids'][i:batch_end] batch_documents = new_memories['documents'][i:batch_end] batch_metadatas = new_memories['metadatas'][i:batch_end] print(f"Migrating batch {i//batch_size + 1} ({len(batch_ids)} memories)...") target_collection.add( documents=batch_documents, metadatas=batch_metadatas, ids=batch_ids ) # Small delay between batches time.sleep(1) print("\nMigration complete!") # Verify migration target_results = target_collection.get() print(f"Verification: {len(target_results['ids'])} total memories in target collection") except Exception as e: print(f"Error during migration: {str(e)}") print("Please ensure both ChromaDB instances are running and accessible") if __name__ == "__main__": # First verify the environment verify_environment() # Example usage: # Local to remote migration migrate_memories( source_type='local', source_config=get_claude_desktop_chroma_path(), target_type='remote', target_config={'host': '16.171.169.46', 'port': 8000} ) # Remote to local migration # migrate_memories( # source_type='remote', # source_config={'host': '16.171.169.46', 'port': 8000}, # target_type='local', # target_config=get_claude_desktop_chroma_path() # )