feat: Qdrant storage support

Cargo.toml

@@ -97,6 +97,7 @@ rustls = { version = "0.23.25" }
 http-body-util = "0.1.3"
 yaml-rust2 = "0.10.1"
 urlencoding = "2.1.3"
+qdrant-client = "1.13.0"
 uuid = { version = "1.16.0", features = ["serde", "v4", "v8"] }
 tokio-stream = "0.1.17"
 async-stream = "0.3.6"

docs/docs/ops/storages.md

@@ -7,10 +7,28 @@ description: CocoIndex Built-in Storages
 
 ## Postgres
 
-`Postgres` exports data to Postgres database (with pgvector extension).
+Exports data to Postgres database (with pgvector extension).
 
 The spec takes the following fields:
 
 *   `database_url` (type: `str`, optional): The URL of the Postgres database to use as the internal storage, e.g. `postgres://cocoindex:cocoindex@localhost/cocoindex`. If unspecified, will use the same database as the [internal storage](/docs/core/basics#internal-storage).
 
 *   `table_name` (type: `str`, optional): The name of the table to store to. If unspecified, will generate a new automatically. We recommend specifying a name explicitly if you want to directly query the table. It can be omitted if you want to use CocoIndex's query handlers to query the table.
+
+## Qdrant
+
+Exports data to a [Qdrant](https://qdrant.tech/) collection.
+
+The spec takes the following fields:
+
+*   `collection_name` (type: `str`, required): The name of the collection to export the data to.
+
+*   `grpc_url` (type: `str`, optional): The [gRPC URL](https://qdrant.tech/documentation/interfaces/#grpc-interface) of the Qdrant instance. Defaults to `http://localhost:6334/`.
+
+*   `api_key` (type: `str`, optional). API key to authenticate requests with.
+
+The field name for the vector embeddings must match the [vector name](https://qdrant.tech/documentation/concepts/vectors/#named-vectors) used when the collection was created.
+
+If no primary key is set during export, a random UUID is used as the Qdrant point ID.
+
+You can find an end-to-end example [here](https://github.com/cocoindex-io/cocoindex/tree/main/examples/text_embedding).

examples/pdf_embedding/main.py

@@ -61,7 +61,7 @@ def pdf_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoinde
 
     doc_embeddings.export(
         "doc_embeddings",
-        cocoindex.storages.Postgres(),
+        cocoindex.storages.Qdrant(collection_name="cocoindex"),
         primary_key_fields=["id"],
         vector_index=[("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
 

examples/text_embedding/README.md

@@ -1,7 +1,34 @@
-Simple example for cocoindex: build embedding index based on local files.
+## Description
 
-## Prerequisite
-[Install Postgres](https://cocoindex.io/docs/getting_started/installation#-install-postgres) if you don't have one.
+Example to build a vector index in Qdrant based on local files.
+
+## Pre-requisites
+
+- [Install Postgres](https://cocoindex.io/docs/getting_started/installation#-install-postgres) if you don't have one.
+
+- Run Qdrant.
+
+```bash
+docker run -d -p 6334:6334 -p 6333:6333 qdrant/qdrant
+```
+
+- [Create a collection](https://qdrant.tech/documentation/concepts/vectors/#named-vectors) to export the embeddings to.
+
+```bash
+curl  -X PUT \
+  'http://localhost:6333/collections/cocoindex' \
+  --header 'Content-Type: application/json' \
+  --data-raw '{
+  "vectors": {
+    "text_embedding": {
+      "size": 384,
+      "distance": "Cosine"
+    }
+  }
+}'
+```
+
+You can view the collections and data with the Qdrant dashboard at <http://localhost:6333/dashboard>.
 
 ## Run
 
@@ -29,13 +56,14 @@ Run:
 python main.py
 ```
 
-## CocoInsight 
+## CocoInsight
+
 CocoInsight is in Early Access now (Free) 😊 You found us! A quick 3 minute video tutorial about CocoInsight: [Watch on YouTube](https://youtu.be/ZnmyoHslBSc?si=pPLXWALztkA710r9).
 
 Run CocoInsight to understand your RAG data pipeline:
 
-```
+```bash
 python main.py cocoindex server -c https://cocoindex.io
 ```
 
-Then open the CocoInsight UI at [https://cocoindex.io/cocoinsight](https://cocoindex.io/cocoinsight).
+Then open the CocoInsight UI at [https://cocoindex.io/cocoinsight](https://cocoindex.io/cocoinsight).

examples/text_embedding/main.py

@@ -2,57 +2,79 @@
 
 import cocoindex
 
+
 def text_to_embedding(text: cocoindex.DataSlice) -> cocoindex.DataSlice:
     """
     Embed the text using a SentenceTransformer model.
     This is a shared logic between indexing and querying, so extract it as a function.
     """
     return text.transform(
         cocoindex.functions.SentenceTransformerEmbed(
-            model="sentence-transformers/all-MiniLM-L6-v2"))
+            model="sentence-transformers/all-MiniLM-L6-v2"
+        )
+    )
+
 
 @cocoindex.flow_def(name="TextEmbedding")
-def text_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope):
+def text_embedding_flow(
+    flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope
+):
     """
     Define an example flow that embeds text into a vector database.
     """
     data_scope["documents"] = flow_builder.add_source(
-        cocoindex.sources.LocalFile(path="markdown_files"))
+        cocoindex.sources.LocalFile(path="markdown_files")
+    )
 
     doc_embeddings = data_scope.add_collector()
 
     with data_scope["documents"].row() as doc:
         doc["chunks"] = doc["content"].transform(
             cocoindex.functions.SplitRecursively(),
-            language="markdown", chunk_size=2000, chunk_overlap=500)
+            language="markdown",
+            chunk_size=2000,
+            chunk_overlap=500,
+        )
 
         with doc["chunks"].row() as chunk:
             chunk["embedding"] = text_to_embedding(chunk["text"])
-            doc_embeddings.collect(filename=doc["filename"], location=chunk["location"],
-                                   text=chunk["text"], embedding=chunk["embedding"])
+            doc_embeddings.collect(
+                id=cocoindex.GeneratedField.UUID,
+                filename=doc["filename"],
+                location=chunk["location"],
+                text=chunk["text"],
+                # 'text_embedding' is the name of the vector we've created the Qdrant collection with.
+                text_embedding=chunk["embedding"],
+            )
 
     doc_embeddings.export(
         "doc_embeddings",
-        cocoindex.storages.Postgres(),
-        primary_key_fields=["filename", "location"],
-        vector_index=[("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
+        cocoindex.storages.Qdrant(
+            collection_name="cocoindex", grpc_url="http://localhost:6334/"
+        ),
+        primary_key_fields=["id"],
+        setup_by_user=True,
+    )
+
 
 query_handler = cocoindex.query.SimpleSemanticsQueryHandler(
     name="SemanticsSearch",
     flow=text_embedding_flow,
     target_name="doc_embeddings",
     query_transform_flow=text_to_embedding,
-    default_similarity_metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)
+    default_similarity_metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY,
+)
+
 
 @cocoindex.main_fn()
 def _run():
     # Run queries in a loop to demonstrate the query capabilities.
     while True:
         try:
             query = input("Enter search query (or Enter to quit): ")
-            if query == '':
+            if query == "":
                 break
-            results, _ = query_handler.search(query, 10)
+            results, _ = query_handler.search(query, 10, "text_embedding")
             print("\nSearch results:")
             for result in results:
                 print(f"[{result.score:.3f}] {result.data['filename']}")
@@ -62,6 +84,7 @@ def _run():
         except KeyboardInterrupt:
             break
 
+
 if __name__ == "__main__":
     load_dotenv(override=True)
     _run()

python/cocoindex/storages.py

@@ -4,12 +4,21 @@
 from . import op
 from . import index
 from .auth_registry import AuthEntryReference
+
 class Postgres(op.StorageSpec):
     """Storage powered by Postgres and pgvector."""
 
     database_url: str | None = None
     table_name: str | None = None
 
+@dataclass
+class Qdrant(op.StorageSpec):
+    """Storage powered by Qdrant - https://qdrant.tech/."""
+
+    collection_name: str
+    grpc_url: str = "http://localhost:6334/"
+    api_key: str | None = None
+
 @dataclass
 class Neo4jConnectionSpec:
     """Connection spec for Neo4j."""

src/base/value.rs

@@ -174,10 +174,7 @@ impl std::fmt::Display for KeyValue {
 }
 
 impl KeyValue {
-    pub fn fields_iter<'a>(
-        &'a self,
-        num_fields: usize,
-    ) -> Result<impl Iterator<Item = &'a KeyValue>> {
+    pub fn fields_iter(&self, num_fields: usize) -> Result<impl Iterator<Item = &KeyValue>> {
         let slice = if num_fields == 1 {
             std::slice::from_ref(self)
         } else {

src/ops/interface.rs

@@ -200,6 +200,7 @@ pub enum ExecutorFactory {
     ExportTarget(Arc<dyn ExportTargetFactory + Send + Sync>),
 }
 
+#[derive(Debug)]
 pub struct VectorMatchQuery {
     pub vector_field_name: String,
     pub vector: Vec<f32>,

src/ops/registration.rs

@@ -13,6 +13,7 @@ fn register_executor_factories(registry: &mut ExecutorFactoryRegistry) -> Result
     functions::extract_by_llm::Factory.register(registry)?;
 
     Arc::new(storages::postgres::Factory::default()).register(registry)?;
+    Arc::new(storages::qdrant::Factory::default()).register(registry)?;
 
     let neo4j_pool = Arc::new(storages::neo4j::GraphPool::default());
     storages::neo4j::RelationshipFactory::new(neo4j_pool).register(registry)?;

src/ops/storages/mod.rs

@@ -1,2 +1,3 @@
 pub mod neo4j;
 pub mod postgres;
+pub mod qdrant;