retrieval/vector_store.py

from typing import List, Dict

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

from config import QDRANT_HOST, QDRANT_PORT, QDRANT_API_KEY, QDRANT_URL, QDRANT_COLLECTION, VECTOR_SIZE

_client = None


def _get_client() -> QdrantClient:
    global _client
    if _client is None:
        if QDRANT_API_KEY and QDRANT_URL:
            _client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
        else:
            _client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
    return _client


def create_collection_if_not_exists():
    client = _get_client()
    existing = [c.name for c in client.get_collections().collections]
    if QDRANT_COLLECTION not in existing:
        client.create_collection(
            collection_name=QDRANT_COLLECTION,
            vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
        )
        print(f"Created collection: {QDRANT_COLLECTION}")
    else:
        print(f"Collection already exists: {QDRANT_COLLECTION}")


def upsert_chunks(chunks: List[Dict], batch_size: int = 100) -> int:
    client = _get_client()
    create_collection_if_not_exists()

    points = [
        PointStruct(
            id=chunk["chunk_id"],
            vector=chunk["embedding"],
            payload={
                "chunk_text": chunk["chunk_text"],
                "source": chunk["source"],
                "metadata": chunk["metadata"],
            },
        )
        for chunk in chunks
        if chunk["embedding"]
    ]

    print(f"Upserting {len(points)} points into '{QDRANT_COLLECTION}' in batches of {batch_size}...")
    for i in range(0, len(points), batch_size):
        batch = points[i: i + batch_size]
        client.upsert(collection_name=QDRANT_COLLECTION, points=batch)
        print(f"  Upserted batch {i // batch_size + 1}/{-(-len(points) // batch_size)}")

    print(f"Done. {len(points)} points upserted.")
    return len(points)


def get_collection_info() -> Dict:
    client = _get_client()
    info = client.get_collection(QDRANT_COLLECTION)
    return {
        "name": QDRANT_COLLECTION,
        "points_count": info.points_count,
        "vector_size": info.config.params.vectors.size,
    }


def delete_collection():
    client = _get_client()
    client.delete_collection(QDRANT_COLLECTION)
    print(f"Deleted collection: {QDRANT_COLLECTION}")