Initial commit: library-desk service extraction from portainer-core

src/clients/qdrant_client.py

@@ -0,0 +1,588 @@
+"""
+Qdrant client wrapper for Library Desk.
+
+Provides async vector operations with:
+- Collection-per-user multi-tenancy
+- Document chunk storage with embeddings
+- Semantic search
+- Similarity queries
+"""
+
+from qdrant_client import QdrantClient
+from qdrant_client.models import (
+    Distance, VectorParams, PointStruct,
+    Filter, FieldCondition, MatchValue
+)
+from typing import List, Dict, Any, Optional
+import uuid
+import logging
+
+from src.core.multi_tenancy import get_qdrant_collection_name
+
+logger = logging.getLogger(__name__)
+
+
+class QdrantClientWrapper:
+    """
+    Qdrant client wrapper with multi-tenancy support.
+
+    Pattern: Collection per user (from qdrant_memory.py)
+    Each user has isolated vector collection for their documents.
+    """
+
+    def __init__(self, url: str, embedding_dim: int = 768):
+        """
+        Initialize Qdrant client.
+
+        Args:
+            url: Qdrant server URL (e.g., "http://qdrant:6333")
+            embedding_dim: Vector embedding dimension (default 768 for nomic-embed-text)
+        """
+        self.client = QdrantClient(url=url)
+        self.embedding_dim = embedding_dim
+        logger.info(f"Initialized Qdrant client: {url}")
+
+    def get_collection_name(self, user: str) -> str:
+        """
+        Get collection name for user.
+
+        Args:
+            user: User identifier
+
+        Returns:
+            Collection name (e.g., "library_desk_jpmschweitzer")
+        """
+        return get_qdrant_collection_name(user)
+
+    async def ensure_collection(self, collection_name: str):
+        """
+        Create collection if it doesn't exist.
+
+        Args:
+            collection_name: Collection name
+        """
+        try:
+            collections = self.client.get_collections()
+            existing = [c.name for c in collections.collections]
+
+            if collection_name not in existing:
+                logger.info(f"Creating Qdrant collection: {collection_name}")
+                self.client.create_collection(
+                    collection_name=collection_name,
+                    vectors_config=VectorParams(
+                        size=self.embedding_dim,
+                        distance=Distance.COSINE
+                    )
+                )
+                logger.info(f"Created collection: {collection_name}")
+        except Exception as e:
+            logger.error(f"Error ensuring collection: {e}", exc_info=True)
+            raise
+
+    async def collection_exists(self, collection_name: str) -> bool:
+        """
+        Check if collection exists.
+
+        Args:
+            collection_name: Collection name
+
+        Returns:
+            True if collection exists
+        """
+        try:
+            collections = self.client.get_collections()
+            existing = [c.name for c in collections.collections]
+            return collection_name in existing
+        except Exception as e:
+            logger.error(f"Error checking collection: {e}", exc_info=True)
+            return False
+
+    async def upsert_document_chunks(
+        self,
+        user: str,
+        doc_id: str,
+        chunks: List[Dict[str, Any]],
+        embeddings: List[List[float]]
+    ) -> int:
+        """
+        Upsert document chunks with embeddings.
+
+        Point structure:
+        {
+            id: "doc_123_chunk_0",
+            vector: [...],
+            payload: {
+                doc_id: "doc_123",
+                chunk_index: 0,
+                content: "text content",
+                metadata: {...}
+            }
+        }
+
+        Args:
+            user: User identifier
+            doc_id: Document ID
+            chunks: List of chunk dictionaries with content
+            embeddings: List of embedding vectors
+
+        Returns:
+            Number of chunks upserted
+
+        Raises:
+            ValueError: If chunks and embeddings length mismatch
+        """
+        if len(chunks) != len(embeddings):
+            raise ValueError(
+                f"Chunks ({len(chunks)}) and embeddings ({len(embeddings)}) length mismatch"
+            )
+
+        collection_name = self.get_collection_name(user)
+        await self.ensure_collection(user)
+
+        points = []
+        for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
+            # Generate deterministic point ID
+            point_id = str(uuid.uuid5(
+                uuid.NAMESPACE_DNS,
+                f"{doc_id}_chunk_{i}"
+            ))
+
+            points.append(PointStruct(
+                id=point_id,
+                vector=embedding,
+                payload={
+                    "doc_id": doc_id,
+                    "chunk_index": i,
+                    "content": chunk.get("content", ""),
+                    "metadata": chunk.get("metadata", {})
+                }
+            ))
+
+        try:
+            self.client.upsert(
+                collection_name=collection_name,
+                points=points
+            )
+            logger.info(f"Upserted {len(points)} chunks for {doc_id}")
+            return len(points)
+        except Exception as e:
+            logger.error(f"Failed to upsert chunks: {e}", exc_info=True)
+            raise
+
+    async def search(
+        self,
+        user: str,
+        query_vector: List[float],
+        limit: int = 10,
+        score_threshold: float = 0.7,
+        filter_dict: Optional[Dict[str, Any]] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        Semantic search in user's collection.
+
+        Args:
+            user: User identifier
+            query_vector: Query embedding vector
+            limit: Maximum results to return
+            score_threshold: Minimum similarity score (0.0-1.0)
+            filter_dict: Optional payload filters
+
+        Returns:
+            List of matches with scores and payloads
+
+        Example:
+            results = await client.search(
+                user="jpmschweitzer",
+                query_vector=[0.1, 0.2, ...],
+                limit=5,
+                filter_dict={"doc_id": "doc_123"}
+            )
+        """
+        collection_name = self.get_collection_name(user)
+
+        # Build filter if provided
+        query_filter = None
+        if filter_dict:
+            conditions = []
+            for key, value in filter_dict.items():
+                conditions.append(
+                    FieldCondition(key=key, match=MatchValue(value=value))
+                )
+            query_filter = Filter(must=conditions)
+
+        try:
+            response = self.client.query_points(
+                collection_name=collection_name,
+                query=query_vector,
+                limit=limit,
+                score_threshold=score_threshold,
+                query_filter=query_filter,
+                with_payload=True
+            )
+
+            return [
+                {
+                    "id": str(point.id),
+                    "score": point.score,
+                    "doc_id": point.payload["doc_id"],
+                    "chunk_index": point.payload["chunk_index"],
+                    "content": point.payload["content"],
+                    "metadata": point.payload.get("metadata", {})
+                }
+                for point in response.points
+            ]
+        except Exception as e:
+            logger.error(f"Search failed: {e}", exc_info=True)
+            return []
+
+    async def delete_document(self, user: str, doc_id: str) -> bool:
+        """
+        Delete all chunks for a document.
+
+        Args:
+            user: User identifier
+            doc_id: Document ID
+
+        Returns:
+            True if deleted successfully
+        """
+        collection_name = self.get_collection_name(user)
+
+        try:
+            self.client.delete(
+                collection_name=collection_name,
+                points_selector=Filter(
+                    must=[
+                        FieldCondition(
+                            key="doc_id",
+                            match=MatchValue(value=doc_id)
+                        )
+                    ]
+                )
+            )
+            logger.info(f"Deleted chunks for {doc_id}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to delete document: {e}", exc_info=True)
+            return False
+
+    async def get_document_chunks(
+        self,
+        user: str,
+        doc_id: str
+    ) -> List[Dict[str, Any]]:
+        """
+        Get all chunks for a document.
+
+        Args:
+            user: User identifier
+            doc_id: Document ID
+
+        Returns:
+            List of chunks with content and metadata
+        """
+        collection_name = self.get_collection_name(user)
+
+        try:
+            # Scroll through points with doc_id filter
+            points, _ = self.client.scroll(
+                collection_name=collection_name,
+                scroll_filter=Filter(
+                    must=[
+                        FieldCondition(
+                            key="doc_id",
+                            match=MatchValue(value=doc_id)
+                        )
+                    ]
+                ),
+                limit=1000,
+                with_payload=True,
+                with_vectors=False
+            )
+
+            return [
+                {
+                    "id": str(point.id),
+                    "chunk_index": point.payload["chunk_index"],
+                    "content": point.payload["content"],
+                    "metadata": point.payload.get("metadata", {})
+                }
+                for point in points
+            ]
+        except Exception as e:
+            logger.error(f"Failed to get document chunks: {e}", exc_info=True)
+            return []
+
+    async def count_documents(self, user: str) -> int:
+        """
+        Count total number of unique documents in user's collection.
+
+        Args:
+            user: User identifier
+
+        Returns:
+            Number of unique documents
+        """
+        collection_name = self.get_collection_name(user)
+
+        try:
+            # Get collection info
+            collection_info = self.client.get_collection(collection_name)
+            # This gives total points, not unique docs
+            # For unique docs, would need to aggregate by doc_id
+            return collection_info.points_count
+        except Exception as e:
+            logger.error(f"Failed to count documents: {e}", exc_info=True)
+            return 0
+
+    async def delete_collection(self, user: str) -> bool:
+        """
+        Delete user's entire collection.
+
+        Warning: This removes all data for the user!
+
+        Args:
+            user: User identifier
+
+        Returns:
+            True if deleted successfully
+        """
+        collection_name = self.get_collection_name(user)
+
+        try:
+            self.client.delete_collection(collection_name)
+            logger.warning(f"Deleted collection: {collection_name}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to delete collection: {e}", exc_info=True)
+            return False
+
+    async def find_similar_chunks(
+        self,
+        user: str,
+        doc_id: str,
+        limit: int = 10
+    ) -> List[Dict[str, Any]]:
+        """
+        Find chunks similar to those in a given document.
+
+        Strategy: Get all chunks from doc, use their vectors to find similar chunks.
+
+        Args:
+            user: User identifier
+            doc_id: Source document ID
+            limit: Maximum results per chunk
+
+        Returns:
+            List of similar chunks from other documents
+        """
+        collection_name = self.get_collection_name(user)
+
+        try:
+            # Get source document chunks with vectors
+            source_points, _ = self.client.scroll(
+                collection_name=collection_name,
+                scroll_filter=Filter(
+                    must=[
+                        FieldCondition(
+                            key="doc_id",
+                            match=MatchValue(value=doc_id)
+                        )
+                    ]
+                ),
+                limit=10,  # Sample first 10 chunks
+                with_payload=True,
+                with_vectors=True
+            )
+
+            if not source_points:
+                return []
+
+            # Search using first chunk's vector
+            # (could aggregate multiple chunks for better results)
+            first_vector = source_points[0].vector
+
+            response = self.client.query_points(
+                collection_name=collection_name,
+                query=first_vector,
+                limit=limit * 2,  # Get more to filter out same doc
+                with_payload=True
+            )
+
+            # Filter out chunks from same document
+            similar = [
+                {
+                    "id": str(point.id),
+                    "score": point.score,
+                    "doc_id": point.payload["doc_id"],
+                    "content": point.payload["content"]
+                }
+                for point in response.points
+                if point.payload["doc_id"] != doc_id
+            ]
+
+            return similar[:limit]
+
+        except Exception as e:
+            logger.error(f"Failed to find similar chunks: {e}", exc_info=True)
+            return []
+
+    async def upsert_vector(
+        self,
+        collection_name: str,
+        vector_id: str,
+        vector: List[float],
+        payload: Dict[str, Any]
+    ) -> bool:
+        """
+        Upsert a single vector point.
+
+        Args:
+            collection_name: Collection name
+            vector_id: Point ID
+            vector: Embedding vector
+            payload: Point payload/metadata
+
+        Returns:
+            True if successful
+        """
+        try:
+            self.client.upsert(
+                collection_name=collection_name,
+                points=[PointStruct(
+                    id=vector_id,
+                    vector=vector,
+                    payload=payload
+                )]
+            )
+            return True
+        except Exception as e:
+            logger.error(f"Failed to upsert vector: {e}", exc_info=True)
+            return False
+
+    async def delete_by_filter(
+        self,
+        collection_name: str,
+        filter_conditions: Dict[str, Any]
+    ) -> int:
+        """
+        Delete points matching filter conditions.
+
+        Args:
+            collection_name: Collection name
+            filter_conditions: Filter conditions (e.g., {"page_id": 5})
+
+        Returns:
+            Number of points deleted (approximation)
+        """
+        try:
+            # Build filter
+            conditions = []
+            for key, value in filter_conditions.items():
+                conditions.append(
+                    FieldCondition(key=key, match=MatchValue(value=value))
+                )
+            query_filter = Filter(must=conditions)
+
+            # Delete points
+            result = self.client.delete(
+                collection_name=collection_name,
+                points_selector=query_filter
+            )
+
+            # Return operation status (Qdrant doesn't return count directly)
+            logger.info(f"Deleted points with filter {filter_conditions}")
+            return 1  # Placeholder, actual count not available from API
+
+        except Exception as e:
+            logger.error(f"Failed to delete by filter: {e}", exc_info=True)
+            return 0
+
+    async def search_vectors(
+        self,
+        collection_name: str,
+        query_vector: List[float],
+        limit: int = 10,
+        score_threshold: float = 0.5,
+        filter_conditions: Optional[Dict[str, Any]] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        Search vectors in collection.
+
+        Args:
+            collection_name: Collection name
+            query_vector: Query embedding vector
+            limit: Maximum results
+            score_threshold: Minimum similarity score
+            filter_conditions: Optional filter conditions
+
+        Returns:
+            List of search results with scores and payloads
+        """
+        # Build filter if provided
+        query_filter = None
+        if filter_conditions:
+            conditions = []
+            for key, value in filter_conditions.items():
+                conditions.append(
+                    FieldCondition(key=key, match=MatchValue(value=value))
+                )
+            query_filter = Filter(must=conditions)
+
+        try:
+            response = self.client.query_points(
+                collection_name=collection_name,
+                query=query_vector,
+                limit=limit,
+                score_threshold=score_threshold,
+                query_filter=query_filter,
+                with_payload=True
+            )
+
+            return [
+                {
+                    "id": str(point.id),
+                    "score": point.score,
+                    "payload": dict(point.payload)
+                }
+                for point in response.points
+            ]
+        except Exception as e:
+            logger.error(f"Search failed: {e}", exc_info=True)
+            return []
+
+    async def list_collections(self) -> List[Dict[str, Any]]:
+        """
+        List all collections with stats.
+
+        Returns:
+            List of collection info dictionaries
+        """
+        try:
+            collections = self.client.get_collections()
+            result = []
+
+            for coll in collections.collections:
+                # Get detailed collection info
+                try:
+                    info = self.client.get_collection(coll.name)
+                    result.append({
+                        "name": coll.name,
+                        "vectors_count": info.vectors_count or 0,
+                        "points_count": info.points_count or 0,
+                        "segments_count": info.segments_count or 0
+                    })
+                except Exception as e:
+                    logger.warning(f"Failed to get info for collection {coll.name}: {e}")
+                    result.append({
+                        "name": coll.name,
+                        "vectors_count": 0,
+                        "points_count": 0,
+                        "segments_count": 0
+                    })
+
+            return result
+
+        except Exception as e:
+            logger.error(f"Failed to list collections: {e}", exc_info=True)
+            return []