""" Milvus Schema Migration Script for Domain-Aware RAG. Adds domain_category and domain_terms fields to existing Milvus collection. This script performs a zero-downtime migration by: 1. Creating a new collection with updated schema 2. Copying existing data to new collection 3. Replacing old collection with new one IMPORTANT: Run this during low-traffic period. Backup data before running. Usage: python scripts/migrate_milvus_schema.py [--dry-run] """ import argparse import sys import time from typing import Any import structlog from pymilvus import ( Collection, CollectionSchema, DataType, FieldSchema, connections, utility, ) from src.core.config import settings logger = structlog.get_logger(__name__) def create_updated_schema(vector_dim: int = 3072) -> CollectionSchema: """Create updated Milvus schema with domain fields. Args: vector_dim: Dimension of embedding vectors (default: 3072 for text-embedding-3-large) Returns: CollectionSchema with domain_category and domain_terms fields added """ fields = [ FieldSchema( name="id", dtype=DataType.INT64, is_primary=True, auto_id=True, ), FieldSchema( name="vector", dtype=DataType.FLOAT_VECTOR, dim=vector_dim, ), FieldSchema( name="record_id", dtype=DataType.VARCHAR, max_length=36, ), FieldSchema( name="document_id", dtype=DataType.VARCHAR, max_length=36, ), FieldSchema( name="filename", dtype=DataType.VARCHAR, max_length=255, ), FieldSchema( name="document_type", dtype=DataType.VARCHAR, max_length=20, ), FieldSchema( name="text_content", dtype=DataType.VARCHAR, max_length=65535, ), FieldSchema( name="metadata", dtype=DataType.JSON, ), # NEW FIELDS for Domain-Aware RAG FieldSchema( name="domain_category", dtype=DataType.VARCHAR, max_length=64, default_value="general", # Default for backward compatibility ), FieldSchema( name="domain_terms", dtype=DataType.ARRAY, element_type=DataType.VARCHAR, max_capacity=50, # Max 50 terms per chunk max_length=100, # Max 100 chars per term ), ] schema = CollectionSchema( fields, description="DSOL document records with embeddings and domain terms", ) return schema def migrate_collection( old_collection_name: str, dry_run: bool = False ) -> None: """Migrate existing collection to new schema with domain fields. Migration steps: 1. Verify old collection exists 2. Create temporary new collection with updated schema 3. Copy all data from old to new collection 4. Create indexes on new collection 5. Drop old collection 6. Rename new collection to original name Args: old_collection_name: Name of existing collection to migrate dry_run: If True, only simulate migration without making changes Raises: ValueError: If old collection doesn't exist Exception: If migration fails """ logger.info( "migration_started", collection=old_collection_name, dry_run=dry_run ) # Connect to Milvus connections.connect( alias="default", host=settings.milvus_host, port=settings.milvus_port, ) # Verify old collection exists if not utility.has_collection(old_collection_name): raise ValueError(f"Collection '{old_collection_name}' does not exist") # Get old collection info old_collection = Collection(old_collection_name) old_collection.load() num_entities = old_collection.num_entities logger.info("old_collection_info", num_entities=num_entities) if dry_run: logger.info("dry_run_mode", message="Would migrate collection", entities=num_entities) return # Create new collection with updated schema new_collection_name = f"{old_collection_name}_migrated" logger.info("creating_new_collection", name=new_collection_name) # Drop if exists (in case of retry) if utility.has_collection(new_collection_name): logger.warning("dropping_existing_temp_collection", name=new_collection_name) utility.drop_collection(new_collection_name) # Get vector dimension from old collection vector_field = None for field in old_collection.schema.fields: if field.dtype == DataType.FLOAT_VECTOR: vector_field = field break if not vector_field: raise ValueError("No vector field found in old collection") vector_dim = vector_field.params.get("dim", 3072) logger.info("detected_vector_dim", dim=vector_dim) # Create new collection new_schema = create_updated_schema(vector_dim) new_collection = Collection( name=new_collection_name, schema=new_schema, ) # Copy data in batches logger.info("copying_data", total_entities=num_entities) batch_size = 1000 copied = 0 try: # Query all data from old collection in batches for offset in range(0, num_entities, batch_size): # Query batch from old collection results = old_collection.query( expr=f"id >= 0", # Get all entities output_fields=["*"], limit=batch_size, offset=offset, ) if not results: break # Prepare data for new collection # Add empty arrays for domain_terms and "general" for domain_category batch_data = [] for entity in results: # Preserve existing fields new_entity = [ entity["vector"], entity["record_id"], entity["document_id"], entity["filename"], entity["document_type"], entity["text_content"], entity["metadata"], "general", # domain_category default [], # domain_terms empty array ] batch_data.append(new_entity) # Transpose for columnar format columnar_data = list(map(list, zip(*batch_data))) # Insert into new collection new_collection.insert(columnar_data) new_collection.flush() copied += len(results) logger.info( "batch_copied", offset=offset, batch_size=len(results), total_copied=copied, progress_pct=round((copied / num_entities) * 100, 1) ) # Small delay to avoid overwhelming Milvus time.sleep(0.1) logger.info("data_copy_complete", total_copied=copied) # Create indexes on new collection logger.info("creating_indexes") # Vector index (HNSW for similarity search) vector_index_params = { "metric_type": "COSINE", "index_type": "HNSW", "params": { "M": 16, "efConstruction": 200, }, } new_collection.create_index( field_name="vector", index_params=vector_index_params, ) # Scalar index for domain_category (for filtering) new_collection.create_index( field_name="domain_category", index_params={"index_type": "INVERTED"}, ) # Array index for domain_terms (for array_contains_any filtering) new_collection.create_index( field_name="domain_terms", index_params={"index_type": "INVERTED"}, ) logger.info("indexes_created") # Load new collection into memory new_collection.load() logger.info("new_collection_loaded") # Verify data integrity new_num_entities = new_collection.num_entities if new_num_entities != num_entities: raise ValueError( f"Data integrity check failed: " f"old={num_entities}, new={new_num_entities}" ) logger.info("data_integrity_verified", entities=new_num_entities) # Release old collection old_collection.release() # Drop old collection logger.info("dropping_old_collection", name=old_collection_name) utility.drop_collection(old_collection_name) # Rename new collection to old name logger.info("renaming_new_collection", old_name=new_collection_name, new_name=old_collection_name) utility.rename_collection(new_collection_name, old_collection_name) logger.info( "migration_complete", collection=old_collection_name, entities_migrated=new_num_entities ) except Exception as e: logger.error( "migration_failed", error=str(e), copied=copied, total=num_entities ) # Clean up new collection on failure if utility.has_collection(new_collection_name): logger.info("cleaning_up_failed_migration", collection=new_collection_name) utility.drop_collection(new_collection_name) raise finally: # Disconnect connections.disconnect("default") def verify_schema(collection_name: str) -> None: """Verify that collection has updated schema with domain fields. Args: collection_name: Name of collection to verify """ logger.info("verifying_schema", collection=collection_name) connections.connect( alias="default", host=settings.milvus_host, port=settings.milvus_port, ) if not utility.has_collection(collection_name): logger.error("collection_not_found", name=collection_name) return collection = Collection(collection_name) # Check for domain fields field_names = [field.name for field in collection.schema.fields] has_category = "domain_category" in field_names has_terms = "domain_terms" in field_names logger.info( "schema_verification", has_domain_category=has_category, has_domain_terms=has_terms, all_fields=field_names, ) if has_category and has_terms: logger.info("schema_verification_passed") else: logger.warning("schema_verification_failed", missing_fields=[ f for f in ["domain_category", "domain_terms"] if f not in field_names ]) connections.disconnect("default") def main(): """Main entry point for migration script.""" parser = argparse.ArgumentParser( description="Migrate Milvus collection to add domain fields" ) parser.add_argument( "--dry-run", action="store_true", help="Simulate migration without making changes" ) parser.add_argument( "--verify-only", action="store_true", help="Only verify schema, don't migrate" ) parser.add_argument( "--collection", default=settings.milvus_collection, help=f"Collection name (default: {settings.milvus_collection})" ) args = parser.parse_args() try: if args.verify_only: verify_schema(args.collection) else: migrate_collection(args.collection, dry_run=args.dry_run) except Exception as e: logger.error("migration_error", error=str(e)) sys.exit(1) if __name__ == "__main__": main()