salim0986/graph-bug-ai

1#!/usr/bin/env python32"""3Initialize Qdrant Cloud cluster with required collections.4Run this script after creating your Qdrant Cloud cluster.5"""6 7import os8import sys9from qdrant_client import QdrantClient10from qdrant_client.models import Distance, VectorParams11 12def initialize_qdrant(url: str, api_key: str = None):13    """Create collections for code search."""14    15    print("🔗 Connecting to Qdrant...")16    17    if api_key:18        client = QdrantClient(url=url, api_key=api_key)19        print("✅ Connected with API key authentication")20    else:21        client = QdrantClient(url=url)22        print("✅ Connected without authentication (local)")23    24    try:25        # Check if collection exists26        collections = client.get_collections()27        existing_names = [c.name for c in collections.collections]28        29        print(f"\n📊 Existing collections: {existing_names if existing_names else 'None'}\n")30        31        collection_name = "repo_code"32        33        if collection_name in existing_names:34            print(f"⚠️  Collection '{collection_name}' already exists!")35            response = input("Do you want to recreate it? (yes/no): ")36            if response.lower() == 'yes':37                client.delete_collection(collection_name)38                print(f"🗑️  Deleted existing collection '{collection_name}'")39            else:40                print("✅ Keeping existing collection")41                return True42        43        # Create collection44        print(f"📦 Creating collection '{collection_name}'...")45        client.create_collection(46            collection_name=collection_name,47            vectors_config=VectorParams(48                size=384,  # all-MiniLM-L6-v2 embedding size49                distance=Distance.COSINE50            )51        )52        print(f"✅ Collection '{collection_name}' created successfully!")53        54        # Show collection info55        info = client.get_collection(collection_name)56        print(f"\n📈 Collection Info:")57        print(f"   Name: {collection_name}")58        print(f"   Vector Size: {info.config.params.vectors.size}")59        print(f"   Distance: {info.config.params.vectors.distance}")60        print(f"   Points Count: {info.points_count}")61        62        print("\n🎉 Qdrant initialization complete!")63        64        return True65        66    except Exception as e:67        print(f"\n❌ Error: {e}")68        return False69 70 71if __name__ == "__main__":72    # Get credentials from environment or command line73    url = os.getenv("QDRANT_URL")74    api_key = os.getenv("QDRANT_API_KEY")75    76    if len(sys.argv) >= 2:77        url = sys.argv[1]78    if len(sys.argv) >= 3:79        api_key = sys.argv[2]80    81    if not url:82        print("❌ Missing Qdrant URL!")83        print("\nUsage:")84        print("  python init_qdrant.py <url> [api_key]")85        print("\nOr set environment variables:")86        print("  export QDRANT_URL='https://xxxxx.aws.cloud.qdrant.io'")87        print("  export QDRANT_API_KEY='your-api-key'  # Optional for cloud")88        print("  python init_qdrant.py")89        sys.exit(1)90    91    print("=" * 60)92    print("🚀 Qdrant Cloud Initialization Script")93    print("=" * 60)94    print(f"URL: {url}")95    print(f"API Key: {'***' + api_key[-8:] if api_key else 'None (local mode)'}")96    print("=" * 60)97    98    success = initialize_qdrant(url, api_key)99    sys.exit(0 if success else 1)100