feat(qdrant): support automatic setup for Qdrant (#577)

* cleanup(qdrant): remove query related logic for qdrant

* refactor(qdrant): prepare vector info ahead of time

* feat(qdrant): support automatic setup

Author: badmonster0

examples/image_search/main.py

@@ -17,8 +17,10 @@
 from transformers import CLIPModel, CLIPProcessor
 
 
-QDRANT_GRPC_URL = os.getenv("QDRANT_GRPC_URL", "http://localhost:6334/")
+QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6334/")
+QDRANT_COLLECTION = "cocoindex_image_search"
 CLIP_MODEL_NAME = "openai/clip-vit-large-patch14"
+CLIP_MODEL_DIMENSION = 768
 
 
 @functools.cache
@@ -40,7 +42,9 @@ def embed_query(text: str) -> list[float]:
 
 
 @cocoindex.op.function(cache=True, behavior_version=1, gpu=True)
-def embed_image(img_bytes: bytes) -> cocoindex.Vector[cocoindex.Float32, Literal[384]]:
+def embed_image(
+    img_bytes: bytes,
+) -> cocoindex.Vector[cocoindex.Float32, Literal[CLIP_MODEL_DIMENSION]]:
     """
     Convert image to embedding using CLIP model.
     """
@@ -56,7 +60,7 @@ def embed_image(img_bytes: bytes) -> cocoindex.Vector[cocoindex.Float32, Literal
 @cocoindex.flow_def(name="ImageObjectEmbedding")
 def image_object_embedding_flow(
     flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope
-):
+) -> None:
     data_scope["images"] = flow_builder.add_source(
         cocoindex.sources.LocalFile(
             path="img", included_patterns=["*.jpg", "*.jpeg", "*.png"], binary=True
@@ -73,14 +77,17 @@ def image_object_embedding_flow(
             filename=img["filename"],
             embedding=img["embedding"],
         )
+
+    qdrant_conn = cocoindex.add_auth_entry(
+        "Qdrant", cocoindex.storages.QdrantConnection(url=QDRANT_URL)
+    )
     img_embeddings.export(
         "img_embeddings",
         cocoindex.storages.Qdrant(
-            collection_name="image_search",
-            grpc_url=QDRANT_GRPC_URL,
+            connection=qdrant_conn,
+            collection_name=QDRANT_COLLECTION,
         ),
         primary_key_fields=["id"],
-        setup_by_user=True,
     )
 
 
@@ -103,7 +110,7 @@ def startup_event():
     load_dotenv()
     cocoindex.init()
     # Initialize Qdrant client
-    app.state.qdrant_client = QdrantClient(url=QDRANT_GRPC_URL, prefer_grpc=True)
+    app.state.qdrant_client = QdrantClient(url=QDRANT_URL, prefer_grpc=True)
     app.state.live_updater = cocoindex.FlowLiveUpdater(image_object_embedding_flow)
     app.state.live_updater.start()
 
@@ -118,7 +125,7 @@ def search(
 
     # Search in Qdrant
     search_results = app.state.qdrant_client.search(
-        collection_name="image_search",
+        collection_name=QDRANT_COLLECTION,
         query_vector=("embedding", query_embedding),
         limit=limit,
     )

examples/text_embedding_qdrant/main.py

@@ -4,8 +4,8 @@
 import cocoindex
 
 # Define Qdrant connection constants
-QDRANT_GRPC_URL = "http://localhost:6334"
-QDRANT_COLLECTION = "cocoindex"
+QDRANT_URL = "http://localhost:6334"
+QDRANT_COLLECTION = "cocoindex_text_embedding"
 
 
 @cocoindex.transform_flow()
@@ -26,7 +26,7 @@ def text_to_embedding(
 @cocoindex.flow_def(name="TextEmbeddingWithQdrant")
 def text_embedding_flow(
     flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope
-):
+) -> None:
     """
     Define an example flow that embeds text into a vector database.
     """
@@ -55,19 +55,22 @@ def text_embedding_flow(
                 text_embedding=chunk["embedding"],
             )
 
+    qdrant_conn = cocoindex.add_auth_entry(
+        "Qdrant", cocoindex.storages.QdrantConnection(url=QDRANT_URL)
+    )
     doc_embeddings.export(
         "doc_embeddings",
         cocoindex.storages.Qdrant(
-            collection_name=QDRANT_COLLECTION, grpc_url=QDRANT_GRPC_URL
+            connection=qdrant_conn,
+            collection_name=QDRANT_COLLECTION,
         ),
         primary_key_fields=["id"],
-        setup_by_user=True,
     )
 
 
-def _main():
+def _main() -> None:
     # Initialize Qdrant client
-    client = QdrantClient(url=QDRANT_GRPC_URL, prefer_grpc=True)
+    client = QdrantClient(url=QDRANT_URL, prefer_grpc=True)
 
     # Run queries in a loop to demonstrate the query capabilities.
     while True:
@@ -87,6 +90,8 @@ def _main():
         for result in search_results:
             score = result.score
             payload = result.payload
+            if payload is None:
+                continue
             print(f"[{score:.3f}] {payload['filename']}")
             print(f"    {payload['text']}")
             print("---")

python/cocoindex/storages.py

@@ -16,13 +16,20 @@ class Postgres(op.StorageSpec):
     table_name: str | None = None
 
 
+@dataclass
+class QdrantConnection:
+    """Connection spec for Qdrant."""
+
+    url: str
+    api_key: str | None = None
+
+
 @dataclass
 class Qdrant(op.StorageSpec):
     """Storage powered by Qdrant - https://qdrant.tech/."""
 
+    connection: AuthEntryReference[QdrantConnection]
     collection_name: str
-    grpc_url: str = "http://localhost:6334/"
-    api_key: str | None = None
 
 
 @dataclass

src/ops/factory_bases.rs

@@ -309,6 +309,12 @@ pub trait StorageFactoryBase: ExportTargetFactory + Send + Sync + 'static {
         Vec<(Self::Key, Self::SetupState)>,
     )>;
 
+    /// Deserialize the setup key from a JSON value.
+    /// You can override this method to provide a custom deserialization logic, e.g. to perform backward compatible deserialization.
+    fn deserialize_setup_key(key: serde_json::Value) -> Result<Self::Key> {
+        Ok(serde_json::from_value(key)?)
+    }
+
     /// Will not be called if it's setup by user.
     /// It returns an error if the target only supports setup by user.
     async fn check_setup_status(
@@ -421,7 +427,7 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
         existing_states: setup::CombinedState<serde_json::Value>,
         auth_registry: &Arc<AuthRegistry>,
     ) -> Result<Box<dyn setup::ResourceSetupStatus>> {
-        let key: T::Key = serde_json::from_value(key.clone())?;
+        let key: T::Key = Self::deserialize_setup_key(key.clone())?;
         let desired_state: Option<T::SetupState> = desired_state
             .map(|v| serde_json::from_value(v.clone()))
             .transpose()?;
@@ -438,12 +444,12 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
     }
 
     fn describe_resource(&self, key: &serde_json::Value) -> Result<String> {
-        let key: T::Key = serde_json::from_value(key.clone())?;
+        let key: T::Key = Self::deserialize_setup_key(key.clone())?;
         StorageFactoryBase::describe_resource(self, &key)
     }
 
     fn normalize_setup_key(&self, key: &serde_json::Value) -> Result<serde_json::Value> {
-        let key: T::Key = serde_json::from_value(key.clone())?;
+        let key: T::Key = Self::deserialize_setup_key(key.clone())?;
         Ok(serde_json::to_value(key)?)
     }
 

src/ops/registration.rs

@@ -3,7 +3,7 @@ use super::{
     storages,
 };
 use anyhow::Result;
-use std::sync::{Arc, LazyLock, RwLock, RwLockReadGuard};
+use std::sync::{LazyLock, RwLock, RwLockReadGuard};
 
 fn register_executor_factories(registry: &mut ExecutorFactoryRegistry) -> Result<()> {
     let reqwest_client = reqwest::Client::new();
@@ -17,7 +17,7 @@ fn register_executor_factories(registry: &mut ExecutorFactoryRegistry) -> Result
     functions::extract_by_llm::Factory.register(registry)?;
 
     storages::postgres::Factory::default().register(registry)?;
-    Arc::new(storages::qdrant::Factory::default()).register(registry)?;
+    storages::qdrant::register(registry)?;
     storages::kuzu::register(registry, reqwest_client)?;
 
     storages::neo4j::Factory::new().register(registry)?;