FEAT: added qdrant

aimakerspace/vectordatabase.py

@@ -52,8 +52,77 @@ class VectorDatabase:
         for text, embedding in zip(list_of_text, embeddings):
             self.insert(text, np.array(embedding))
         return self
+import hashlib
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import PointStruct
+class QdrantDatabase:
+    def __init__(self, qdrant_client: QdrantClient, collection_name: str, embedding_model=None):
+        self.qdrant_client = qdrant_client
+        self.collection_name = collection_name
+        self.embedding_model = embedding_model or EmbeddingModel()
+        self.vectors = defaultdict(np.array)  # Still keeps a local copy if needed
+
+    def string_to_int_id(self, s: str) -> int:
+        return int(hashlib.sha256(s.encode('utf-8')).hexdigest(), 16) % (10**8)
+    
+    def insert(self, key: str, vector: np.array) -> None:
+        
+        point_id = self.string_to_int_id(key)
+        # Insert vector into Qdrant
+        payload = {"text": key}  # Storing the key (text) as payload
+        point = PointStruct(
+            id=point_id,
+            vector={"default": vector.tolist()},  # Use the vector name defined in the collection
+            payload=payload
+        )
+
+        # Insert the vector into Qdrant with the associated document
+        self.qdrant_client.upsert(
+            collection_name=self.collection_name,
+            points=[point]  # Qdrant expects a list of PointStruct
+        )
+
+    def search(
+        self,
+        query_vector: np.array,
+        k: int,
+        distance_measure: Callable = None,
+    ) -> List[Tuple[str, float]]:
+        # Perform search in Qdrant
+        print(query_vector)
+        if isinstance(query_vector, list):
+            query_vector = np.array(query_vector)
+    
+        search_results = self.qdrant_client.search(
+            collection_name=self.collection_name,
+            query_vector={"name": "default", "vector": query_vector.tolist()},# Convert numpy array to list
+            limit=k
+        )
 
+        # Extract and return results
+        return [(result.payload['text'], result.score) for result in search_results]
 
+    def search_by_text(
+        self,
+        query_text: str,
+        k: int,
+        distance_measure: Callable = None,
+        return_as_text: bool = False,
+    ) -> List[Tuple[str, float]]:
+        query_vector = self.embedding_model.get_embedding(query_text)
+        results = self.search(query_vector, k, distance_measure)
+        return [result[0] for result in results] if return_as_text else results
+
+    def retrieve_from_key(self, key: str) -> np.array:
+        # Retrieve from local cache
+        return self.vectors.get(key, None)
+
+    async def abuild_from_list(self, list_of_text: List[str]) -> "QdrantDatabase":
+        embeddings = await self.embedding_model.async_get_embeddings(list_of_text)
+        for text, embedding in zip(list_of_text, embeddings):
+            self.insert(text, np.array(embedding))
+        return self
+    
 if __name__ == "__main__":
     list_of_text = [
         "I like to eat broccoli and bananas.",

app.py

@@ -8,7 +8,7 @@ from aimakerspace.openai_utils.prompts import (
     AssistantRolePrompt,
 )
 from aimakerspace.openai_utils.embedding import EmbeddingModel
-from aimakerspace.vectordatabase import VectorDatabase
+from aimakerspace.vectordatabase import VectorDatabase, QdrantDatabase
 from aimakerspace.openai_utils.chatmodel import ChatOpenAI
 import chainlit as cl
 import fitz 
@@ -80,6 +80,7 @@ def process_text_file(file: AskFileResponse):
     return texts
 
 
+
 @cl.on_chat_start
 async def on_chat_start():
     files = None
@@ -112,19 +113,40 @@ async def on_chat_start():
 
     # decide if to use the dict vector store of the Qdrant vector store
 
-    use_qdrant = False
-
+    use_qdrant = True
+    from qdrant_client import QdrantClient
+    from qdrant_client.http.models import VectorParams, Distance
     # Create a dict vector store
     if use_qdrant:
-        msg = cl.Message(
-            content="Sorry, qgrant not implemented yet", disable_human_feedback=True
+        embedding_model = EmbeddingModel()
+        qdrant_client = QdrantClient(
+            url='https://6b3eac94-adfe-42cb-98f8-9f068538243c.europe-west3-0.gcp.cloud.qdrant.io:6333',  # Replace with your cluster URL
+            api_key='YrnApyEfdNAt41N7WkcZwjhjKqiIQQbXHBtzk_04guNyRLa83J0hOw'    # Replace with your API key
         )
-        await msg.send()
-        raise NotImplemented()
+        vectors_config = {
+            "default": VectorParams(size=1536, distance="Cosine")  # Adjust size as per your model's output
+        }
+        if not qdrant_client.collection_exists("my_collection"):
+            qdrant_client.create_collection(
+                collection_name="my_collection",
+                vectors_config=vectors_config
+            )
+
+        vector_db = QdrantDatabase(
+            qdrant_client=qdrant_client,
+            collection_name="my_collection",
+            embedding_model=embedding_model  # Replace with your embedding model instance
+        )
+        vector_db = await vector_db.abuild_from_list(texts)
+
     else:
         vector_db = VectorDatabase()
         vector_db = await vector_db.abuild_from_list(texts)
         
+    msg = cl.Message(
+        content=f"The Vector store has been created", disable_human_feedback=True
+    )
+    await msg.send()
     chat_openai = ChatOpenAI()
 
     # Create a chain