from qdrant_client import QdrantClient, models

import os

from typing import List, Dict, Any

QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")

QDRANT_PORT = int(os.getenv("QDRANT_PORT", 6333))

QDRANT_COLLECTION_NAME = "menu_items"

VECTOR_SIZE = 3072

client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)

def initialize_qdrant_collection(recreate_collection: bool = False):

"""Ensures the Qdrant collection exists with the correct configuration.

If recreate_collection is True, it will delete and recreate the collection."""

if recreate_collection and client.collection_exists(collection_name=QDRANT_COLLECTION_NAME):

client.delete_collection(collection_name=QDRANT_COLLECTION_NAME)

print(f"Qdrant collection '{QDRANT_COLLECTION_NAME}' deleted for recreation.")

if not client.collection_exists(collection_name=QDRANT_COLLECTION_NAME):

client.create_collection(

collection_name=QDRANT_COLLECTION_NAME,

vectors_config=models.VectorParams(size=VECTOR_SIZE, distance=models.Distance.COSINE),

)

print(f"Qdrant collection '{QDRANT_COLLECTION_NAME}' created with vector size {VECTOR_SIZE}")

else:

print(f"Qdrant collection '{QDRANT_COLLECTION_NAME}' already exists (expected size {VECTOR_SIZE}).")

def upload_vectors_to_qdrant(

vectors: List[List[float]],

payloads: List[Dict[str, Any]],

ids: List[int]

):

"""Uploads vectors and their associated payloads to Qdrant."""

client.upsert(

collection_name=QDRANT_COLLECTION_NAME,

wait=True,

points=models.Batch(

ids=ids,

vectors=vectors,

payloads=payloads,

)

)

print(f"Uploaded {len(ids)} points to Qdrant collection '{QDRANT_COLLECTION_NAME}'.")