ai-robot-core/ai-service/app/core/qdrant_client.py

"""
Qdrant client for AI Service.
[AC-AISVC-10] Vector database client with tenant-isolated collection management.
"""

import logging
from typing import Any

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

from app.core.config import get_settings

logger = logging.getLogger(__name__)

settings = get_settings()


class QdrantClient:
    """
    [AC-AISVC-10] Qdrant client with tenant-isolated collection management.
    Collection naming: kb_{tenantId} for tenant isolation.
    """

    def __init__(self):
        self._client: AsyncQdrantClient | None = None
        self._collection_prefix = settings.qdrant_collection_prefix
        self._vector_size = settings.qdrant_vector_size

    async def get_client(self) -> AsyncQdrantClient:
        """Get or create Qdrant client instance."""
        if self._client is None:
            self._client = AsyncQdrantClient(url=settings.qdrant_url)
            logger.info(f"[AC-AISVC-10] Qdrant client initialized: {settings.qdrant_url}")
        return self._client

    async def close(self) -> None:
        """Close Qdrant client connection."""
        if self._client:
            await self._client.close()
            self._client = None
            logger.info("Qdrant client connection closed")

    def get_collection_name(self, tenant_id: str) -> str:
        """
        [AC-AISVC-10] Get collection name for a tenant.
        Naming convention: kb_{tenantId}
        """
        return f"{self._collection_prefix}{tenant_id}"

    async def ensure_collection_exists(self, tenant_id: str) -> bool:
        """
        [AC-AISVC-10] Ensure collection exists for tenant.
        Note: MVP uses pre-provisioned collections, this is for development/testing.
        """
        client = await self.get_client()
        collection_name = self.get_collection_name(tenant_id)

        try:
            collections = await client.get_collections()
            exists = any(c.name == collection_name for c in collections.collections)

            if not exists:
                await client.create_collection(
                    collection_name=collection_name,
                    vectors_config=VectorParams(
                        size=self._vector_size,
                        distance=Distance.COSINE,
                    ),
                )
                logger.info(
                    f"[AC-AISVC-10] Created collection: {collection_name} for tenant={tenant_id}"
                )
            return True
        except Exception as e:
            logger.error(f"[AC-AISVC-10] Error ensuring collection: {e}")
            return False

    async def upsert_vectors(
        self,
        tenant_id: str,
        points: list[PointStruct],
    ) -> bool:
        """
        [AC-AISVC-10] Upsert vectors into tenant's collection.
        """
        client = await self.get_client()
        collection_name = self.get_collection_name(tenant_id)

        try:
            await client.upsert(
                collection_name=collection_name,
                points=points,
            )
            logger.info(
                f"[AC-AISVC-10] Upserted {len(points)} vectors for tenant={tenant_id}"
            )
            return True
        except Exception as e:
            logger.error(f"[AC-AISVC-10] Error upserting vectors: {e}")
            return False

    async def search(
        self,
        tenant_id: str,
        query_vector: list[float],
        limit: int = 5,
        score_threshold: float | None = None,
    ) -> list[dict[str, Any]]:
        """
        [AC-AISVC-10] Search vectors in tenant's collection.
        Returns results with score >= score_threshold if specified.
        """
        client = await self.get_client()
        collection_name = self.get_collection_name(tenant_id)

        try:
            results = await client.search(
                collection_name=collection_name,
                query_vector=query_vector,
                limit=limit,
            )

            hits = [
                {
                    "id": str(result.id),
                    "score": result.score,
                    "payload": result.payload or {},
                }
                for result in results
                if score_threshold is None or result.score >= score_threshold
            ]

            logger.info(
                f"[AC-AISVC-10] Search returned {len(hits)} results for tenant={tenant_id}"
            )
            return hits
        except Exception as e:
            logger.error(f"[AC-AISVC-10] Error searching vectors: {e}")
            return []

    async def delete_collection(self, tenant_id: str) -> bool:
        """
        [AC-AISVC-10] Delete tenant's collection.
        Used when tenant is removed.
        """
        client = await self.get_client()
        collection_name = self.get_collection_name(tenant_id)

        try:
            await client.delete_collection(collection_name=collection_name)
            logger.info(f"[AC-AISVC-10] Deleted collection: {collection_name}")
            return True
        except Exception as e:
            logger.error(f"[AC-AISVC-10] Error deleting collection: {e}")
            return False


_qdrant_client: QdrantClient | None = None


async def get_qdrant_client() -> QdrantClient:
    """Get or create Qdrant client instance."""
    global _qdrant_client
    if _qdrant_client is None:
        _qdrant_client = QdrantClient()
    return _qdrant_client


async def close_qdrant_client() -> None:
    """Close Qdrant client connection."""
    global _qdrant_client
    if _qdrant_client:
        await _qdrant_client.close()
        _qdrant_client = None