scripts/seed_qdrant.py

import asyncio
import os
import urllib.request
import tempfile
from qdrant_client import AsyncQdrantClient
from qdrant_client.models import PointStruct, VectorParams, Distance
from dotenv import load_dotenv

load_dotenv(os.path.join(os.path.dirname(__file__), '..', '.env'))

QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
COLLECTION_NAME = "currency_standards"

async def seed_qdrant():
    print(f"Connecting to Qdrant at {QDRANT_URL}...")
    client = AsyncQdrantClient(url=QDRANT_URL)
    
    # Check if collection exists
    collections = await client.get_collections()
    if not any(c.name == COLLECTION_NAME for c in collections.collections):
        print(f"Creating collection '{COLLECTION_NAME}'...")
        await client.create_collection(
            collection_name=COLLECTION_NAME,
            vectors_config=VectorParams(size=768, distance=Distance.COSINE)
        )
    
    print("Downloading genuine Indian Currency (INR) references...")
    
    # Genuine reference notes from Wikimedia Commons
    inr_notes = [
        {
            "id": 1,
            "denomination": 500,
            "url": "https://upload.wikimedia.org/wikipedia/commons/2/2e/India_new_500_INR%2C_MG_series%2C_2016%2C_obverse.jpg"
        },
        {
            "id": 2,
            "denomination": 2000,
            "url": "https://upload.wikimedia.org/wikipedia/commons/0/07/India_new_2000_INR%2C_MG_series%2C_2016%2C_obverse.jpg"
        }
    ]
    
    # Import the Vision Service to use GroundingDINO and DINOv2
    # We do this here so we only load models if this script is executed
    print("Loading AI Vision Models (this may take a moment)...")
    import sys
    sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
    from ml_services.vision_ai import vision_service
    
    points = []
    
    for note in inr_notes:
        print(f"Processing INR {note['denomination']}...")
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file:
            urllib.request.urlretrieve(note["url"], tmp_file.name)
            
            # Use the actual Vision Service to crop and extract the true DINO embedding
            embedding, bbox, serial = await asyncio.to_thread(vision_service._process_image_sync, tmp_file.name)
            
            points.append(
                PointStruct(
                    id=note["id"],
                    vector=embedding,
                    payload={
                        "currency_type": "INR",
                        "denomination": note["denomination"],
                        "description": "Genuine reference note",
                        "source": "Wikimedia Commons"
                    }
                )
            )
            os.remove(tmp_file.name)
        
    print("Uploading DINOv2 reference vectors to Qdrant...")
    await client.upsert(
        collection_name=COLLECTION_NAME,
        points=points
    )
    
    print("Successfully seeded Qdrant with authentic DINOv2 embeddings of genuine Indian Currency!")

if __name__ == "__main__":
    asyncio.run(seed_qdrant())