Initial commit

app/services/vector_store.py

@@ -0,0 +1,153 @@
+from qdrant_client import QdrantClient, models
+from qdrant_client.http.models import (
+    Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue
+)
+from app.config import settings
+from typing import List, Dict, Any, Optional
+import logging
+import uuid
+
+logger = logging.getLogger(__name__)
+
+_qdrant_client: QdrantClient = None
+
+
+def get_qdrant_client() -> QdrantClient:
+    global _qdrant_client
+    if _qdrant_client is None:
+        kwargs = {"url": settings.qdrant_url}
+        if settings.qdrant_api_key:
+            kwargs["api_key"] = settings.qdrant_api_key
+        _qdrant_client = QdrantClient(**kwargs)
+    return _qdrant_client
+
+
+class VectorStoreService:
+    VECTOR_SIZE = 1536  # text-embedding-3-small
+
+    def __init__(self):
+        self.client = get_qdrant_client()
+
+    def create_collection(self, collection_name: str) -> bool:
+        """Create a new collection for a chatbot"""
+        try:
+            self.client.create_collection(
+                collection_name=collection_name,
+                vectors_config=VectorParams(
+                    size=self.VECTOR_SIZE,
+                    distance=Distance.COSINE,
+                ),
+            )
+            logger.info(f"Created collection: {collection_name}")
+            return True
+        except Exception as e:
+            if "already exists" in str(e).lower():
+                return True
+            logger.error(f"Error creating collection {collection_name}: {e}")
+            raise
+
+    def delete_collection(self, collection_name: str) -> bool:
+        """Delete a chatbot's collection"""
+        try:
+            self.client.delete_collection(collection_name=collection_name)
+            logger.info(f"Deleted collection: {collection_name}")
+            return True
+        except Exception as e:
+            logger.error(f"Error deleting collection {collection_name}: {e}")
+            return False
+
+    def collection_exists(self, collection_name: str) -> bool:
+        try:
+            self.client.get_collection(collection_name)
+            return True
+        except Exception:
+            return False
+
+    def upsert_vectors(
+        self,
+        collection_name: str,
+        vectors: List[List[float]],
+        payloads: List[Dict[str, Any]],
+        ids: Optional[List[str]] = None,
+    ) -> bool:
+        """Upsert vectors into collection"""
+        if ids is None:
+            ids = [str(uuid.uuid4()) for _ in vectors]
+
+        points = [
+            PointStruct(
+                id=idx,
+                vector=vector,
+                payload=payload,
+            )
+            for idx, vector, payload in zip(ids, vectors, payloads)
+        ]
+
+        try:
+            self.client.upsert(
+                collection_name=collection_name,
+                points=points,
+            )
+            return True
+        except Exception as e:
+            logger.error(f"Error upserting vectors: {e}")
+            raise
+
+    def search(
+        self,
+        collection_name: str,
+        query_vector: List[float],
+        limit: int = 5,
+        score_threshold: float = 0.3,
+    ) -> List[Dict[str, Any]]:
+        """Search for similar vectors"""
+        try:
+            results = self.client.search(
+                collection_name=collection_name,
+                query_vector=query_vector,
+                limit=limit,
+                score_threshold=score_threshold,
+            )
+            return [
+                {
+                    "id": str(r.id),
+                    "score": r.score,
+                    "payload": r.payload,
+                }
+                for r in results
+            ]
+        except Exception as e:
+            logger.error(f"Error searching vectors: {e}")
+            return []
+
+    def delete_by_document_id(self, collection_name: str, document_id: str) -> bool:
+        """Delete all vectors for a document"""
+        try:
+            self.client.delete(
+                collection_name=collection_name,
+                points_selector=models.FilterSelector(
+                    filter=Filter(
+                        must=[
+                            FieldCondition(
+                                key="document_id",
+                                match=MatchValue(value=document_id),
+                            )
+                        ]
+                    )
+                ),
+            )
+            return True
+        except Exception as e:
+            logger.error(f"Error deleting document vectors: {e}")
+            return False
+
+    def count_vectors(self, collection_name: str) -> int:
+        """Count vectors in a collection"""
+        try:
+            result = self.client.count(collection_name=collection_name)
+            return result.count
+        except Exception:
+            return 0
+
+
+vector_store = VectorStoreService()