scripts/migrate_qdrant_collection.py

#!/usr/bin/env python3
"""
Production-safe Qdrant embedding-model migration.

Run from inside the Docker backend container:
    docker exec chatbot_backend_dev uv run python scripts/migrate_qdrant_collection.py \\
        --target smart_chatbot_kb_v2

Steps executed:
    1. Inspect source collection dimension vs configured QDRANT_VECTOR_SIZE
    2. If dimensions already match → nothing to do, exit 0
    3. Create target collection with the configured QDRANT_VECTOR_SIZE
    4. Re-embed all tenant KB entries into the target collection (source unchanged)
    5. Run health probe on target collection
    6. Print cutover instructions

Flags:
    --target    (required) Name of the new Qdrant collection to create and populate.
    --dry-run   Inspect dimensions and print the migration plan without writing anything.
"""

import argparse
import sys
from contextlib import contextmanager

from sqlalchemy import select

from app.models.database import SessionLocal
from app.models.smart_models import KnowledgeBase
from app.services.embedding_service import EmbeddingService
from app.services.kb_reindex import reindex_tenant
from app.services.qdrant_vector_store import QdrantVectorStore
from app.utils.config import config_manager


def run_migration(
    source_store: QdrantVectorStore,
    target_store: QdrantVectorStore,
    embedding_service: EmbeddingService,
    db_factory,
    dry_run: bool = False,
) -> None:
    """Core migration logic. Extracted for testability.

    Args:
        source_store:      QdrantVectorStore pointing at the current live collection.
        target_store:      QdrantVectorStore pointing at the new collection to populate.
        embedding_service: EmbeddingService configured for the target embedding model.
        db_factory:        Callable context-manager returning a SQLAlchemy Session.
        dry_run:           If True, print the plan but write nothing.
    """
    if source_store.collection_name == target_store.collection_name:
        print(
            f"Error: --target must be different from the source collection name "
            f"({source_store.collection_name}). Aborting."
        )
        sys.exit(1)

    source_dim = source_store.get_collection_dimension()
    target_dim = target_store.vector_size

    print(f"Source collection : {source_store.collection_name}")
    print(f"  Actual dimension: {source_dim if source_dim is not None else 'does not exist'}")
    print(f"Target collection : {target_store.collection_name}")
    print(f"  Target dimension: {target_dim}")

    if source_dim is not None and source_dim == target_dim:
        print("\nNo migration needed — source and target dimensions already match.")
        sys.exit(0)

    if dry_run:
        mismatch = f"{source_dim} → {target_dim}" if source_dim else f"(new) → {target_dim}"
        print(f"\nDry run: would migrate {source_store.collection_name} ({mismatch})"
              f" → {target_store.collection_name}")
        print("No writes performed.")
        return

    # Step 1: Create target collection (or validate existing one)
    existing_target_dim = target_store.get_collection_dimension()
    if existing_target_dim is not None and existing_target_dim != target_dim:
        print(
            f"Error: target collection '{target_store.collection_name}' already exists "
            f"with dim={existing_target_dim}, expected dim={target_dim}. "
            f"Delete it first or choose a different --target name."
        )
        sys.exit(1)
    target_store._ensure_collection()
    print(f"\n✓ Target collection ready: {target_store.collection_name} (dim={target_dim})")

    # Step 2: Backfill all tenant KB entries
    with db_factory() as db:
        tenant_ids = [
            str(tid)
            for tid in db.scalars(
                select(KnowledgeBase.tenant_id).distinct()
            ).all()
        ]

    total_entries = 0
    for tenant_id in tenant_ids:
        with db_factory() as db:
            count = reindex_tenant(tenant_id, embedding_service, target_store, db)
        total_entries += count
        print(f"  Tenant {tenant_id}: {count} entries reindexed")

    print(f"\n✓ Backfill complete: {total_entries} entries across {len(tenant_ids)} tenants")

    # Step 3: Health probe on target collection
    target_store.health_check()
    print("✓ Health probe passed on target collection")

    # Step 4: Cutover instructions
    print("\n─── CUTOVER INSTRUCTIONS ────────────────────────────────────────────────")
    print(f"  Update your .env:  QDRANT_COLLECTION_NAME={target_store.collection_name}")
    print("  Restart backend:   docker compose -f docker-compose.dev.yml restart smart-chatbot")
    print("")
    print("  After verifying the new collection works in production:")
    print(f"  Retire old collection: set QDRANT_COLLECTION_NAME back, call delete via client")
    print(f"  Or:  from qdrant_client import QdrantClient")
    print(f"       QdrantClient(url='...').delete_collection('{source_store.collection_name}')")
    print("─────────────────────────────────────────────────────────────────────────")
    print("\n⚠  Note: KB entries written during migration may be missing from the target")
    print("   collection. After cutover, call POST /api/v1/kb/reindex per tenant to")
    print("   catch any missed entries.")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Migrate Qdrant collection to a new embedding dimension."
    )
    parser.add_argument(
        "--target",
        required=True,
        help="Target collection name (new collection to create and populate).",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Inspect dimensions and print the plan without writing anything.",
    )
    args = parser.parse_args()

    config = config_manager.get_config()

    source_store = QdrantVectorStore(
        url=config.qdrant_url,
        collection_name=config.qdrant_collection_name,
        vector_size=config.qdrant_vector_size,
        timeout=config.qdrant_timeout,
    )
    target_store = QdrantVectorStore(
        url=config.qdrant_url,
        collection_name=args.target,
        vector_size=config.qdrant_vector_size,
        timeout=config.qdrant_timeout,
    )
    embedding_service = EmbeddingService(
        config.embedding_model,
        api_key=config.gemini_embedding_api_key,
    )

    @contextmanager
    def session_factory():
        with SessionLocal() as session:
            yield session

    run_migration(
        source_store=source_store,
        target_store=target_store,
        embedding_service=embedding_service,
        db_factory=session_factory,
        dry_run=args.dry_run,
    )


if __name__ == "__main__":
    main()