Add support for Qdrant VectorDB

docs/griptape-framework/drivers/vector-store-drivers.md

@@ -406,3 +406,112 @@ vector_store_driver.upsert_text_artifacts(
 result = vector_store_driver.query("What is griptape?")
 print(result)
 ```
+
+### Qdrant
+
+!!! info
+    This driver requires the `drivers-vector-qdrant` [extra](../index.md#extras).
+
+The QdrantVectorStoreDriver supports the [Qdrant vector database](https://qdrant.tech/).
+
+Here is an example of how the driver can be used to query information in a Qdrant collection:
+
+```python
+import os
+import logging
+from transformers import AutoTokenizer
+from sentence_transformers import SentenceTransformer
+from griptape.drivers import QdrantVectorStoreDriver, HuggingFaceHubEmbeddingDriver
+from griptape.tokenizers import HuggingFaceTokenizer
+from griptape.loaders import PdfLoader
+from griptape.chunkers import TextChunker
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+
+# Setting the models
+embedding_model = "sentence-transformers/all-MiniLM-L6-v2"
+qdrant_model = SentenceTransformer("all-MiniLM-L6-v2")
+file_name = "linux_bible.pdf"
+HUGGINGFACE_TOKEN = os.environ["HUGGINGFACE_HUB_ACCESS_TOKEN"]
+
+
+# Using HuggingFaceTokenizer
+def create_tokenizer(embedding_model):
+    tokenizer = HuggingFaceTokenizer(max_output_tokens=1024, tokenizer=AutoTokenizer.from_pretrained(embedding_model))
+    return tokenizer
+
+
+def create_embedding_driver(embedding_model, tokenizer):
+    embedding_driver = HuggingFaceHubEmbeddingDriver(
+        api_token=HUGGINGFACE_TOKEN, model=embedding_model, tokenizer=tokenizer
+    )
+    return embedding_driver
+
+
+# Instantiating QdrantVectorStoreDriver
+def create_vector_store_driver(url, collection_name, embedding_driver):
+    vector_store_driver = QdrantVectorStoreDriver(
+        url=url, collection_name=collection_name, content_payload_key="content", embedding_driver=embedding_driver
+    )
+    return vector_store_driver
+
+
+# Opening the file
+def load_pdf(file_name, tokenizer):
+    with open(file_name, "rb") as f:
+        loader = PdfLoader(tokenizer=tokenizer, chunker=TextChunker(tokenizer=tokenizer, max_tokens=1024)).load(
+            f.read()
+        )
+    return loader
+
+
+def main():
+    tokenizer = create_tokenizer(embedding_model)
+    embedding_driver = create_embedding_driver(embedding_model, tokenizer)
+    vector_store_driver = create_vector_store_driver(
+        url="http://localhost:6333", collection_name="linux_bible", embedding_driver=embedding_driver
+    )
+
+    # Loading the data
+    loader = load_pdf(file_name, tokenizer=tokenizer)
+
+    # Generate metadata for each chunk (example metadata)
+    metadata = [{"source": file_name, "page": i + 1} for i in range(len(loader))]
+
+    # Upserting the vector
+    try:
+        for i, l in enumerate(loader):
+            content = str(l)
+            meta = metadata[i] if metadata else None
+            vector_store_driver.upsert_vector(embedding_driver.try_embed_chunk(content), meta=meta, content=content)
+        logging.info("Successfully upserted vectors with metadata.")
+    except Exception as e:
+        logging.error(f"Error during upsert_vector: {e}")
+
+    # Querying the data
+    query_string = "Who created linux?"
+    query_results = vector_store_driver.query(query_string, count=6, include_vectors=True)
+    for result in query_results:
+       print(f"ID: {result.id}, Score: {result.score}, Vector: {result.vector}, Metadata: {result.meta}")
+
+    # Retrieving single entries
+    single_entry = vector_store_driver.load_entry(vector_id="00499354-9362-49c7-97b6-0220b9de84e7")
+    if single_entry:
+        print(f"Vector ID: {single_entry.id}")
+        print(f"Vector: {single_entry.vector}")
+        print(f"Metadata: {single_entry.meta}")
+    else:
+        print("Vector with ID 00499354-9362-49c7-97b6-0220b9de84e7 was not found.")
+
+    # Retrieving multiple entries
+    multiple_entries = vector_store_driver.load_entries(ids=["707a004c-23bc-476c-a097-e527977777f3"])
+    print(multiple_entries)
+
+    # Deleting the vector
+    vector_store_driver.delete_vector(vector_id="00499354-9362-49c7-97b6-0220b9de84e7")
+
+if __name__ == "__main__":
+    main()
+
+```

griptape/drivers/__init__.py

@@ -39,6 +39,7 @@
 from .vector.pgvector_vector_store_driver import PgVectorVectorStoreDriver
 from .vector.azure_mongodb_vector_store_driver import AzureMongoDbVectorStoreDriver
 from .vector.dummy_vector_store_driver import DummyVectorStoreDriver
+from .vector.qdrant_vector_store_driver import QdrantVectorStoreDriver
 
 from .sql.base_sql_driver import BaseSqlDriver
 from .sql.amazon_redshift_sql_driver import AmazonRedshiftSqlDriver
@@ -144,6 +145,7 @@
     "OpenSearchVectorStoreDriver",
     "AmazonOpenSearchVectorStoreDriver",
     "PgVectorVectorStoreDriver",
+    "QdrantVectorStoreDriver",
     "DummyVectorStoreDriver",
     "BaseSqlDriver",
     "AmazonRedshiftSqlDriver",

griptape/drivers/vector/qdrant_vector_store_driver.py

@@ -0,0 +1,210 @@
+from __future__ import annotations
+from typing import Optional
+from attrs import define, field
+from griptape.drivers import BaseVectorStoreDriver
+from griptape.utils import import_optional_dependency
+import uuid
+import logging
+
+VECTOR_NAME = None
+DEFAULT_DISTANCE = "COSINE"
+CONTENT_PAYLOAD_KEY = "data"
+
+
+@define
+class QdrantVectorStoreDriver(BaseVectorStoreDriver):
+    """
+    Attributes:
+        location: An optional location for the Qdrant client. If set to ':memory:', an in-memory client is used.
+        url: An optional Qdrant API URL.
+        host: An optional Qdrant host.
+        path: Persistence path for QdrantLocal. Default: None
+        port: The port number for the Qdrant client. Defaults to 6333.
+        grpc_port: The gRPC port number for the Qdrant client. Defaults to 6334.
+        prefer_grpc: A boolean indicating whether to prefer gRPC over HTTP. Defaults to False.
+        force_disable_check_same_thread: For QdrantLocal, force disable check_same_thread. Default: False Only use this if you can guarantee that you can resolve the thread safety outside QdrantClient.
+        timeout: Timeout for REST and gRPC API requests. Default: 5 seconds for REST and unlimited for gRPC
+        api_key: API key for authentication in Qdrant Cloud. Defaults to False
+        https: If true - use HTTPS(SSL) protocol. Default: None
+        prefix: Add prefix to the REST URL path. Example: service/v1 will result in Example: service/v1 will result in http://localhost:6333/service/v1/{qdrant-endpoint} for REST API. Defaults to None
+        distance: The distance metric to be used for the vectors. Defaults to 'COSINE'.
+        collection_name: The name of the Qdrant collection.
+        vector_name: An optional name for the vectors.
+        content_payload_key: The key for the content payload in the metadata. Defaults to 'data'.
+    """
+
+    location: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
+    url: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
+    host: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
+    path: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
+    port: int = field(default=6333, kw_only=True, metadata={"serializable": True})
+    grpc_port: int = field(default=6334, kw_only=True, metadata={"serializable": True})
+    prefer_grpc: bool = field(default=False, kw_only=True, metadata={"serializable": True})
+    api_key: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
+    https: bool = field(default=None, kw_only=True, metadata={"serializable": True})
+    prefix: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
+    force_disable_check_same_thread: Optional[bool] = field(
+        default=False, kw_only=True, metadata={"serializable": True}
+    )
+    timeout: Optional[int] = field(default=5, kw_only=True, metadata={"serializable": True})
+    distance: str = field(default=DEFAULT_DISTANCE, kw_only=True, metadata={"serializable": True})
+    collection_name: str = field(kw_only=True, metadata={"serializable": True})
+    vector_name: Optional[str] = VECTOR_NAME
+    content_payload_key: str = field(default=CONTENT_PAYLOAD_KEY, kw_only=True, metadata={"serializable": True})
+
+    def __attrs_post_init__(self) -> None:
+        self.client = import_optional_dependency("qdrant_client").QdrantClient(
+            location=self.location,
+            url=self.url,
+            host=self.host,
+            path=self.path,
+            port=self.port,
+            prefer_grpc=self.prefer_grpc,
+            grpc_port=self.grpc_port,
+            api_key=self.api_key,
+            https=self.https,
+            prefix=self.prefix,
+            force_disable_check_same_thread=self.force_disable_check_same_thread,
+            timeout=self.timeout,
+        )
+
+    def delete_vector(self, vector_id: str) -> None:
+        """
+        Delete a vector from the Qdrant collection based on its ID.
+
+        Parameters:
+            vector_id (str | id): ID of the vector to delete.
+        """
+        deletion_response = self.client.delete(
+            collection_name=self.collection_name,
+            points_selector=import_optional_dependency("qdrant_client.http.models").PointIdsList(points=[vector_id]),
+        )
+        if deletion_response.status == import_optional_dependency("qdrant_client.http.models").UpdateStatus.COMPLETED:
+            logging.info(f"ID {vector_id} is successfully deleted")
+
+    def query(
+        self,
+        query: str,
+        count: Optional[int] = None,
+        namespace: Optional[str] = None,
+        include_vectors: bool = False,
+        **kwargs,
+    ) -> list[BaseVectorStoreDriver.Entry]:
+        """
+        Query the Qdrant collection based on a query vector.
+
+        Parameters:
+            query (str): Query string.
+            count (Optional[int]): Optional number of results to return.
+            namespace (Optional[str]): Optional namespace of the vectors.
+            include_vectors (bool): Whether to include vectors in the results.
+
+        Returns:
+            list[BaseVectorStoreDriver.Entry]: List of Entry objects.
+        """
+        query_vector = self.embedding_driver.embed_string(query)
+
+        # Create a search request
+        results = self.client.search(collection_name=self.collection_name, query_vector=query_vector, limit=count)
+
+        # Convert results to QueryResult objects
+        query_results = [
+            BaseVectorStoreDriver.Entry(
+                id=result.id,
+                vector=result.vector if include_vectors else [],
+                score=result.score,
+                meta={k: v for k, v in result.payload.items() if k not in ["_score", "_tensor_facets"]},
+            )
+            for result in results
+        ]
+        return query_results
+
+    def upsert_vector(
+        self,
+        vector: list[float],
+        vector_id: Optional[str] = None,
+        namespace: Optional[str] = None,
+        meta: Optional[dict] = None,
+        content: Optional[str] = None,
+        **kwargs,
+    ) -> str:
+        """
+        Upsert vectors into the Qdrant collection.
+
+        Parameters:
+            vector (list[float]): The vector to be upserted.
+            vector_id (Optional[str]): Optional vector ID.
+            namespace (Optional[str]): Optional namespace for the vector.
+            meta (Optional[dict]): Optional dictionary containing metadata.
+            content (Optional[str]): The text content to be included in the payload.
+
+        Returns:
+            str: The ID of the upserted vector.
+        """
+
+        if vector_id is None:
+            vector_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, str(vector)))
+
+        if meta is None:
+            meta = {}
+
+        if content:
+            meta[self.content_payload_key] = content
+
+        points = import_optional_dependency("qdrant_client.http.models").Batch(
+            ids=[vector_id], vectors=[vector], payloads=[meta] if meta else None
+        )
+
+        self.client.upsert(collection_name=self.collection_name, points=points)
+        return vector_id
+
+    def load_entry(self, vector_id: str, namespace: Optional[str] = None) -> Optional[BaseVectorStoreDriver.Entry]:
+        """
+        Load a vector entry from the Qdrant collection based on its ID.
+
+        Parameters:
+            vector_id (str): ID of the vector to load.
+            namespace (str, optional): Optional namespace of the vector.
+
+        Returns:
+            Optional[BaseVectorStoreDriver.Entry]: Vector entry if found, else None.
+        """
+        results = self.client.retrieve(collection_name=self.collection_name, ids=[vector_id])
+        if results:
+            entry = results[0]
+            return BaseVectorStoreDriver.Entry(
+                id=entry.id,
+                vector=entry.vector,
+                meta={k: v for k, v in entry.payload.items() if k not in ["_score", "_tensor_facets"]},
+            )
+        else:
+            return None
+
+    def load_entries(self, namespace: Optional[str] = None, **kwargs) -> list[BaseVectorStoreDriver.Entry]:
+        """
+        Load vector entries from the Qdrant collection.
+
+        Parameters:
+            namespace: Optional namespace of the vectors.
+
+        Returns:
+            List of points.
+        """
+
+        results = self.client.retrieve(
+            collection_name=self.collection_name,
+            ids=kwargs.get("ids", []),
+            with_payload=kwargs.get("with_payload", True),
+            with_vectors=kwargs.get("with_vectors", True),
+        )
+        if not results:
+            logging.error("An error occurred or no results found.")
+            return []
+        return [
+            BaseVectorStoreDriver.Entry(
+                id=entry.id,
+                vector=entry.vector if kwargs.get("with_vectors", True) else [],
+                meta={k: v for k, v in entry.payload.items() if k not in ["_score", "_tensor_facets"]},
+            )
+            for entry in results
+        ]