rename_collections.py

"""Rename collections to meaningful names."""
import sqlite3
import sys

def rename_collections(collection_mappings):
    """Rename collections in SQLite.
    
    Args:
        collection_mappings: Dict of {collection_id: new_name}
    """
    
    print("\n" + "=" * 80)
    print("🔄 Collection Rename Tool")
    print("=" * 80)
    
    conn = sqlite3.connect('chroma_db/chroma.sqlite3')
    cursor = conn.cursor()
    
    # Show current collections
    print("\n📚 Current Collections:")
    cursor.execute('SELECT id, name FROM collections ORDER BY name')
    current = cursor.fetchall()
    for i, (coll_id, name) in enumerate(current, 1):
        print(f"  {i}. {name} (ID: {coll_id[:8]}...)")
    
    if not collection_mappings:
        print("\n⚠️  No mappings provided")
        conn.close()
        return False
    
    # Apply renames
    print("\n🔄 Applying renames...")
    print("-" * 80)
    
    updated = 0
    for coll_id, new_name in collection_mappings.items():
        try:
            cursor.execute("""
                UPDATE collections 
                SET name = ? 
                WHERE id = ?
            """, (new_name, coll_id))
            
            if cursor.rowcount > 0:
                print(f"✅ {coll_id[:8]}... -> {new_name}")
                updated += 1
            else:
                print(f"❌ Collection not found: {coll_id[:8]}...")
        except Exception as e:
            print(f"❌ Error updating {coll_id[:8]}...: {e}")
    
    # Commit changes
    print("\n💾 Committing changes...")
    try:
        conn.commit()
        print(f"✅ Updated {updated} collection(s)")
    except Exception as e:
        print(f"❌ Error committing: {e}")
        conn.close()
        return False
    
    # Show new names
    print("\n📚 Updated Collections:")
    cursor.execute('SELECT id, name FROM collections ORDER BY name')
    updated_list = cursor.fetchall()
    for i, (coll_id, name) in enumerate(updated_list, 1):
        print(f"  {i}. {name}")
    
    conn.close()
    return True


def main():
    """Main entry point - update collection names here."""
    
    # ✏️ CONFIGURE YOUR COLLECTION NAMES HERE
    collection_mappings = {
        # UUID -> New Name
        '35506af1-bd91-4ae5-a461-b6aa0878d531': 'chunking_embedding_llm',
        '441127d3-4a66-425e-b80a-385e2c58bbce': 'dataset_collection_2',
        '65dcdf88-843d-4978-8527-8fb4ec6cece8': 'dataset_collection_3',
        'c3270053-8004-4f4d-bcb8-32e0cb692d57': 'dataset_collection_4',
    }
    
    success = rename_collections(collection_mappings)
    
    print("\n" + "=" * 80)
    if success:
        print("✅ RENAME COMPLETE!")
        print("\nRestart Streamlit to see updated collection names:")
        print("  streamlit run streamlit_app.py")
    else:
        print("❌ RENAME FAILED")
    print("=" * 80 + "\n")
    
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()