#!/usr/bin/env python3
"""
Migration Script: Add description_colbert Named Vector to ColBERT Collection
This script migrates the existing card_framework_components_colbert collection
to add a second named vector (description_colbert) for the feedback loop.
Qdrant does NOT support adding new named vectors to existing collections,
so we must:
1. Create a new collection with both named vectors
2. Copy all existing points (with placeholder for description_colbert)
3. Rename collections (old -> backup, new -> primary)
Usage:
uv run python scripts/migrate_colbert_collection_for_feedback.py [--dry-run]
"""
import os
import sys
from typing import List
from dotenv import load_dotenv
# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
load_dotenv()
# Configuration
OLD_COLLECTION = "card_framework_components_colbert"
NEW_COLLECTION = "card_framework_components_colbert_v2"
BACKUP_COLLECTION = "card_framework_components_colbert_backup"
COLBERT_DIM = 128
BATCH_SIZE = 100
def get_qdrant_client():
"""Get Qdrant client using .env credentials."""
from qdrant_client import QdrantClient
url = os.getenv("QDRANT_URL")
key = os.getenv("QDRANT_KEY")
if not url or not key:
raise ValueError("QDRANT_URL and QDRANT_KEY must be set in .env")
return QdrantClient(url=url, api_key=key, prefer_grpc=True)
def check_collection_exists(client, name: str) -> bool:
"""Check if a collection exists."""
collections = client.get_collections()
return name in [c.name for c in collections.collections]
def get_collection_info(client, name: str):
"""Get collection info."""
return client.get_collection(name)
def create_new_collection_with_both_vectors(client, dry_run: bool = False):
"""Create the new collection with both named vectors."""
from qdrant_client.models import (
Distance,
MultiVectorComparator,
MultiVectorConfig,
VectorParams,
)
print(f"\n📦 Creating new collection: {NEW_COLLECTION}")
print(f" Named vectors: colbert, description_colbert")
print(f" Dimension: {COLBERT_DIM}")
print(f" Multi-vector comparator: MAX_SIM")
if dry_run:
print(" [DRY RUN] Would create collection")
return
# Check if new collection already exists
if check_collection_exists(client, NEW_COLLECTION):
print(f" ⚠️ Collection {NEW_COLLECTION} already exists - deleting first")
client.delete_collection(NEW_COLLECTION)
client.create_collection(
collection_name=NEW_COLLECTION,
vectors_config={
"colbert": VectorParams(
size=COLBERT_DIM,
distance=Distance.COSINE,
multivector_config=MultiVectorConfig(
comparator=MultiVectorComparator.MAX_SIM
),
),
"description_colbert": VectorParams(
size=COLBERT_DIM,
distance=Distance.COSINE,
multivector_config=MultiVectorConfig(
comparator=MultiVectorComparator.MAX_SIM
),
),
},
)
print(f" ✅ Created {NEW_COLLECTION}")
# Create payload indexes for filtering
print(f"\n📑 Creating payload indexes...")
from qdrant_client.models import PayloadSchemaType
client.create_payload_index(
collection_name=NEW_COLLECTION,
field_name="type",
field_schema=PayloadSchemaType.KEYWORD,
)
print(f" ✅ Created index on 'type' field")
client.create_payload_index(
collection_name=NEW_COLLECTION,
field_name="feedback",
field_schema=PayloadSchemaType.KEYWORD,
)
print(f" ✅ Created index on 'feedback' field")
def migrate_points(client, dry_run: bool = False):
"""Migrate all points from old to new collection."""
from qdrant_client.models import PointStruct
# Get total count
old_info = get_collection_info(client, OLD_COLLECTION)
total_points = old_info.points_count
print(
f"\n📋 Migrating {total_points} points from {OLD_COLLECTION} to {NEW_COLLECTION}"
)
if total_points == 0:
print(" ⚠️ No points to migrate")
return
if dry_run:
print(f" [DRY RUN] Would migrate {total_points} points")
return
# Scroll through all points in batches
offset = None
migrated = 0
while True:
# Scroll to get points
results, next_offset = client.scroll(
collection_name=OLD_COLLECTION,
limit=BATCH_SIZE,
offset=offset,
with_payload=True,
with_vectors=True,
)
if not results:
break
# Convert points for new collection
new_points = []
for point in results:
# Get existing colbert vectors
colbert_vectors = point.vector
if isinstance(colbert_vectors, dict):
colbert_vectors = colbert_vectors.get("colbert", [])
# Create placeholder for description_colbert
# Use a single zero vector (will be overwritten when patterns are stored)
placeholder_vector = [[0.0] * COLBERT_DIM]
new_point = PointStruct(
id=point.id,
vector={
"colbert": colbert_vectors,
"description_colbert": placeholder_vector,
},
payload=point.payload,
)
new_points.append(new_point)
# Upsert batch
client.upsert(
collection_name=NEW_COLLECTION,
points=new_points,
)
migrated += len(new_points)
print(f" ✅ Migrated {migrated}/{total_points} points")
if next_offset is None:
break
offset = next_offset
print(f"\n ✅ Migration complete: {migrated} points")
def swap_collections(client, dry_run: bool = False):
"""Rename collections: old -> backup, new -> primary."""
print(f"\n🔄 Swapping collections:")
print(f" {OLD_COLLECTION} -> {BACKUP_COLLECTION}")
print(f" {NEW_COLLECTION} -> {OLD_COLLECTION}")
if dry_run:
print(" [DRY RUN] Would swap collections")
return
# Delete backup if it exists
if check_collection_exists(client, BACKUP_COLLECTION):
print(f" ⚠️ Deleting existing backup: {BACKUP_COLLECTION}")
client.delete_collection(BACKUP_COLLECTION)
# Rename old -> backup
# Note: Qdrant doesn't have rename, so we need to recreate
# For now, just delete old and rename new
# Actually, Qdrant Cloud supports collection aliases which is the better approach
# Let's check if aliases are available
try:
# Try to use aliases (more elegant)
from qdrant_client.models import (
AliasOperations,
CreateAliasOperation,
DeleteAliasOperation,
)
# Update alias to point to new collection
client.update_collection_aliases(
change_aliases_operations=[
(
DeleteAliasOperation(alias_name=OLD_COLLECTION)
if check_alias_exists(client, OLD_COLLECTION)
else None
),
CreateAliasOperation(
alias_name=OLD_COLLECTION,
collection_name=NEW_COLLECTION,
),
]
)
print(f" ✅ Created alias {OLD_COLLECTION} -> {NEW_COLLECTION}")
except Exception as e:
# Aliases not available or failed, use delete/rename approach
print(f" ℹ️ Aliases not available ({e}), using delete/create approach")
# Delete old collection
print(f" 🗑️ Deleting old collection: {OLD_COLLECTION}")
client.delete_collection(OLD_COLLECTION)
# Rename new collection (by recreating with same name)
# This is destructive - we'll lose the new collection name
# Better approach: just update the code to use NEW_COLLECTION
print(f" ℹ️ New collection available as: {NEW_COLLECTION}")
print(f" ⚠️ Update code to use: {NEW_COLLECTION}")
def check_alias_exists(client, alias_name: str) -> bool:
"""Check if an alias exists."""
try:
aliases = client.get_collection_aliases(alias_name)
return bool(aliases.aliases)
except Exception:
return False
def verify_migration(client):
"""Verify the migration was successful."""
print(f"\n🔍 Verifying migration:")
# Check new collection
if not check_collection_exists(client, NEW_COLLECTION):
print(f" ❌ New collection {NEW_COLLECTION} does not exist!")
return False
new_info = get_collection_info(client, NEW_COLLECTION)
print(f" New collection points: {new_info.points_count}")
# Check vectors config
vectors_config = new_info.config.params.vectors
if isinstance(vectors_config, dict):
print(f" Named vectors: {list(vectors_config.keys())}")
if "colbert" in vectors_config and "description_colbert" in vectors_config:
print(" ✅ Both named vectors present")
else:
print(" ❌ Missing expected named vectors")
return False
# Sample a point to verify structure
results, _ = client.scroll(
collection_name=NEW_COLLECTION,
limit=1,
with_vectors=True,
)
if results:
point = results[0]
if isinstance(point.vector, dict):
print(f" Sample point vectors: {list(point.vector.keys())}")
print(" ✅ Point structure correct")
else:
print(" ❌ Point vector structure incorrect")
return False
return True
def main():
"""Run the migration."""
import argparse
parser = argparse.ArgumentParser(
description="Migrate ColBERT collection for feedback loop"
)
parser.add_argument(
"--dry-run",
action="store_true",
help="Show what would be done without making changes",
)
parser.add_argument(
"--skip-swap",
action="store_true",
help="Don't swap collections, keep new collection with _v2 suffix",
)
args = parser.parse_args()
print("=" * 60)
print("Migration: Add description_colbert to ColBERT Collection")
print("=" * 60)
if args.dry_run:
print("\n🔶 DRY RUN MODE - No changes will be made")
client = get_qdrant_client()
# Check old collection exists
if not check_collection_exists(client, OLD_COLLECTION):
print(f"\n❌ Source collection {OLD_COLLECTION} does not exist!")
return 1
old_info = get_collection_info(client, OLD_COLLECTION)
print(f"\n📊 Source collection: {OLD_COLLECTION}")
print(f" Points: {old_info.points_count}")
vectors_config = old_info.config.params.vectors
if isinstance(vectors_config, dict):
print(f" Named vectors: {list(vectors_config.keys())}")
# Step 1: Create new collection
create_new_collection_with_both_vectors(client, args.dry_run)
# Step 2: Migrate points
if not args.dry_run:
migrate_points(client, args.dry_run)
# Step 3: Verify migration
if not args.dry_run:
if not verify_migration(client):
print("\n❌ Migration verification failed!")
return 1
# Step 4: Swap collections (optional)
if not args.skip_swap:
swap_collections(client, args.dry_run)
else:
print(f"\n⏭️ Skipping collection swap")
print(f" New collection available as: {NEW_COLLECTION}")
print(f" Update COLLECTION_NAME in code to use new collection")
print("\n" + "=" * 60)
print("✅ Migration complete!")
print("=" * 60)
if args.skip_swap or args.dry_run:
print(f"\nNext steps:")
print(
f"1. Update COLLECTION_NAME in gchat/feedback_loop.py to: {NEW_COLLECTION}"
)
print(
f"2. Update COLLECTION_NAME in gchat/smart_card_builder.py to: {NEW_COLLECTION}"
)
print(f"3. Run tests to verify functionality")
return 0
if __name__ == "__main__":
sys.exit(main())
