tatlock/src/core/qdrant.py

"""
Qdrant client wrapper for memory vector storage.

Provides async operations for storing and retrieving memory embeddings:
- Collection management (per-user collections)
- Memory upsert/search/delete
- Filtering by memory type

Adapted from library-desk patterns.
"""
from typing import Any
from uuid import uuid4, uuid5, NAMESPACE_DNS

from qdrant_client import QdrantClient
from qdrant_client.http import models as qdrant_models

from .config import config
from .logging_config import get_logger
from .multi_tenancy import get_memory_collection_name

logger = get_logger(__name__)


class MemoryQdrantClient:
    """
    Qdrant client wrapper for memory storage.

    Manages per-user collections with the pattern: memories_{user}
    Stores memory embeddings with metadata (type, content, timestamps).

    Usage:
        client = MemoryQdrantClient()
        await client.ensure_collection("jpmschweitzer")
        await client.upsert_memory(
            user="jpmschweitzer",
            memory_id="mem_123",
            vector=[0.1, 0.2, ...],
            payload={"type": "fact", "content": "User prefers dark mode"}
        )
    """

    def __init__(
        self,
        url: str | None = None,
        embedding_dim: int | None = None,
    ):
        """
        Initialize Qdrant client.

        Args:
            url: Qdrant server URL (defaults to config.qdrant_url)
            embedding_dim: Vector dimension (defaults to config.QDRANT_EMBEDDING_DIM)
        """
        self.url = url or config.qdrant_url
        self.embedding_dim = embedding_dim or config.QDRANT_EMBEDDING_DIM
        self._client = QdrantClient(url=self.url)

        logger.info(
            "qdrant_client_initialized",
            url=self.url,
            embedding_dim=self.embedding_dim,
        )

    def close(self) -> None:
        """Close Qdrant client."""
        if self._client is not None:
            self._client.close()

    async def ensure_collection(self, user: str) -> bool:
        """
        Ensure collection exists for user, create if not.

        Args:
            user: User identifier

        Returns:
            True if collection exists or was created successfully

        Example:
            >>> await client.ensure_collection("jpmschweitzer")
            True
        """
        collection_name = get_memory_collection_name(user)

        try:
            # Check if collection exists
            collections = self._client.get_collections()
            existing = [c.name for c in collections.collections]

            if collection_name in existing:
                logger.debug(
                    "qdrant_collection_exists",
                    collection=collection_name,
                )
                return True

            # Create collection with cosine distance
            self._client.create_collection(
                collection_name=collection_name,
                vectors_config=qdrant_models.VectorParams(
                    size=self.embedding_dim,
                    distance=qdrant_models.Distance.COSINE,
                ),
            )

            logger.info(
                "qdrant_collection_created",
                collection=collection_name,
                embedding_dim=self.embedding_dim,
            )
            return True

        except Exception as e:
            logger.error(
                "qdrant_ensure_collection_failed",
                collection=collection_name,
                error=str(e),
            )
            return False

    async def upsert_memory(
        self,
        user: str,
        memory_id: str | None,
        vector: list[float],
        payload: dict[str, Any],
    ) -> str | None:
        """
        Upsert a memory point.

        Args:
            user: User identifier
            memory_id: Memory ID (generated if None)
            vector: Embedding vector
            payload: Memory metadata (should include 'type', 'content', etc.)

        Returns:
            Memory ID if successful, None on failure

        Example:
            >>> memory_id = await client.upsert_memory(
            ...     user="jpmschweitzer",
            ...     memory_id=None,
            ...     vector=[0.1, 0.2, ...],
            ...     payload={
            ...         "type": "fact",
            ...         "content": "User prefers dark mode",
            ...         "created_at": "2024-01-01T00:00:00Z"
            ...     }
            ... )
        """
        collection_name = get_memory_collection_name(user)

        # Generate deterministic UUID from memory_id (or random if not provided)
        # Qdrant requires UUID or integer IDs, not arbitrary strings
        if memory_id:
            # Deterministic UUID from string - same memory_id = same UUID
            point_id = str(uuid5(NAMESPACE_DNS, f"{user}:{memory_id}"))
        else:
            point_id = str(uuid4())
            memory_id = point_id  # Use UUID as the memory_id too

        try:
            # Ensure collection exists
            await self.ensure_collection(user)

            # Create point (store original memory_id in payload for reference)
            payload["memory_id"] = memory_id
            point = qdrant_models.PointStruct(
                id=point_id,
                vector=vector,
                payload=payload,
            )

            # Upsert
            self._client.upsert(
                collection_name=collection_name,
                points=[point],
            )

            logger.debug(
                "qdrant_memory_upserted",
                collection=collection_name,
                memory_id=memory_id,
                memory_type=payload.get("type"),
            )
            return memory_id

        except Exception as e:
            logger.error(
                "qdrant_upsert_memory_failed",
                collection=collection_name,
                memory_id=memory_id,
                error=str(e),
            )
            return None

    async def search_memories(
        self,
        user: str,
        query_vector: list[float],
        limit: int = 10,
        memory_type: str | None = None,
        score_threshold: float = 0.5,
    ) -> list[dict[str, Any]]:
        """
        Search memories by vector similarity.

        Args:
            user: User identifier
            query_vector: Query embedding vector
            limit: Maximum results
            memory_type: Filter by memory type (e.g., "fact", "preference", "profile")
            score_threshold: Minimum similarity score (0-1)

        Returns:
            List of matching memories with scores

        Example:
            >>> memories = await client.search_memories(
            ...     user="jpmschweitzer",
            ...     query_vector=[0.1, 0.2, ...],
            ...     limit=5,
            ...     memory_type="fact"
            ... )
            >>> memories[0]
            {"id": "mem_123", "score": 0.89, "type": "fact", "content": "..."}
        """
        collection_name = get_memory_collection_name(user)

        try:
            # Build filter if memory_type specified
            query_filter = None
            if memory_type:
                query_filter = qdrant_models.Filter(
                    must=[
                        qdrant_models.FieldCondition(
                            key="type",
                            match=qdrant_models.MatchValue(value=memory_type),
                        )
                    ]
                )

            # Search using new Query API (qdrant-client >= 1.10)
            results = self._client.query_points(
                collection_name=collection_name,
                query=query_vector,
                limit=limit,
                query_filter=query_filter,
                score_threshold=score_threshold,
            ).points

            # Format results
            memories = []
            for hit in results:
                memory = {
                    "id": hit.id,
                    "score": hit.score,
                    **hit.payload,
                }
                memories.append(memory)

            logger.debug(
                "qdrant_search_memories",
                collection=collection_name,
                results_count=len(memories),
                memory_type=memory_type,
            )
            return memories

        except Exception as e:
            logger.error(
                "qdrant_search_memories_failed",
                collection=collection_name,
                error=str(e),
            )
            return []

    async def get_memory(self, user: str, memory_id: str) -> dict[str, Any] | None:
        """
        Get a specific memory by ID.

        Args:
            user: User identifier
            memory_id: Memory ID

        Returns:
            Memory data or None if not found
        """
        collection_name = get_memory_collection_name(user)
        # Convert memory_id to UUID point_id
        point_id = str(uuid5(NAMESPACE_DNS, f"{user}:{memory_id}"))

        try:
            points = self._client.retrieve(
                collection_name=collection_name,
                ids=[point_id],
            )

            if not points:
                return None

            point = points[0]
            return {
                "id": point.id,
                **point.payload,
            }

        except Exception as e:
            logger.error(
                "qdrant_get_memory_failed",
                collection=collection_name,
                memory_id=memory_id,
                error=str(e),
            )
            return None

    async def delete_memory(self, user: str, memory_id: str) -> bool:
        """
        Delete a memory by ID.

        Args:
            user: User identifier
            memory_id: Memory ID to delete

        Returns:
            True if deleted successfully, False otherwise

        Example:
            >>> await client.delete_memory("jpmschweitzer", "mem_123")
            True
        """
        collection_name = get_memory_collection_name(user)
        # Convert memory_id to UUID point_id
        point_id = str(uuid5(NAMESPACE_DNS, f"{user}:{memory_id}"))

        try:
            self._client.delete(
                collection_name=collection_name,
                points_selector=qdrant_models.PointIdsList(
                    points=[point_id],
                ),
            )

            logger.debug(
                "qdrant_memory_deleted",
                collection=collection_name,
                memory_id=memory_id,
            )
            return True

        except Exception as e:
            logger.error(
                "qdrant_delete_memory_failed",
                collection=collection_name,
                memory_id=memory_id,
                error=str(e),
            )
            return False

    async def delete_memories_by_type(self, user: str, memory_type: str) -> int:
        """
        Delete all memories of a specific type.

        Args:
            user: User identifier
            memory_type: Type of memories to delete

        Returns:
            Number of memories deleted (approximate)
        """
        collection_name = get_memory_collection_name(user)

        try:
            # Delete by filter
            self._client.delete(
                collection_name=collection_name,
                points_selector=qdrant_models.FilterSelector(
                    filter=qdrant_models.Filter(
                        must=[
                            qdrant_models.FieldCondition(
                                key="type",
                                match=qdrant_models.MatchValue(value=memory_type),
                            )
                        ]
                    )
                ),
            )

            logger.info(
                "qdrant_memories_deleted_by_type",
                collection=collection_name,
                memory_type=memory_type,
            )
            return -1  # Qdrant doesn't return count for filter deletes

        except Exception as e:
            logger.error(
                "qdrant_delete_memories_by_type_failed",
                collection=collection_name,
                memory_type=memory_type,
                error=str(e),
            )
            return 0

    async def count_memories(self, user: str) -> int:
        """
        Count total memories for a user.

        Args:
            user: User identifier

        Returns:
            Number of memories in user's collection
        """
        collection_name = get_memory_collection_name(user)

        try:
            info = self._client.get_collection(collection_name)
            return info.points_count

        except Exception as e:
            logger.error(
                "qdrant_count_memories_failed",
                collection=collection_name,
                error=str(e),
            )
            return 0

    async def health_check(self) -> bool:
        """
        Check if Qdrant server is reachable.

        Returns:
            True if healthy, False otherwise
        """
        try:
            self._client.get_collections()
            return True
        except Exception as e:
            logger.error("qdrant_health_check_failed", error=str(e))
            return False


# Global client instance (lazy initialization)
_qdrant_client: MemoryQdrantClient | None = None


def get_qdrant_client() -> MemoryQdrantClient:
    """
    Get global Qdrant client instance.

    Returns:
        MemoryQdrantClient instance
    """
    global _qdrant_client
    if _qdrant_client is None:
        _qdrant_client = MemoryQdrantClient()
    return _qdrant_client