Skip to content

Commit 9dd6eae

Browse files
authored
chore: make connection optional for qdrant spec, s/grpc_url/url/ (#580)
1 parent 058223f commit 9dd6eae

File tree

4 files changed

+30
-31
lines changed

4 files changed

+30
-31
lines changed

examples/image_search/main.py

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,14 @@
1212
from fastapi.middleware.cors import CORSMiddleware
1313
from fastapi.staticfiles import StaticFiles
1414
from qdrant_client import QdrantClient
15+
from typing import Any
1516

1617
from PIL import Image
1718
from transformers import CLIPModel, CLIPProcessor
1819

1920

2021
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6334/")
21-
QDRANT_COLLECTION = "cocoindex_image_search"
22+
QDRANT_COLLECTION = "ImageSearch"
2223
CLIP_MODEL_NAME = "openai/clip-vit-large-patch14"
2324
CLIP_MODEL_DIMENSION = 768
2425

@@ -78,15 +79,9 @@ def image_object_embedding_flow(
7879
embedding=img["embedding"],
7980
)
8081

81-
qdrant_conn = cocoindex.add_auth_entry(
82-
"Qdrant", cocoindex.storages.QdrantConnection(url=QDRANT_URL)
83-
)
8482
img_embeddings.export(
8583
"img_embeddings",
86-
cocoindex.storages.Qdrant(
87-
connection=qdrant_conn,
88-
collection_name=QDRANT_COLLECTION,
89-
),
84+
cocoindex.storages.Qdrant(collection_name=QDRANT_COLLECTION),
9085
primary_key_fields=["id"],
9186
)
9287

@@ -106,7 +101,7 @@ def image_object_embedding_flow(
106101

107102
# --- CocoIndex initialization on startup ---
108103
@app.on_event("startup")
109-
def startup_event():
104+
def startup_event() -> None:
110105
load_dotenv()
111106
cocoindex.init()
112107
# Initialize Qdrant client
@@ -119,7 +114,7 @@ def startup_event():
119114
def search(
120115
q: str = Query(..., description="Search query"),
121116
limit: int = Query(5, description="Number of results"),
122-
):
117+
) -> Any:
123118
# Get the embedding for the query
124119
query_embedding = embed_query(q)
125120

examples/text_embedding_qdrant/main.py

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
from dotenv import load_dotenv
22
from qdrant_client import QdrantClient
3-
from qdrant_client.http.models import Filter, FieldCondition, MatchValue
43
import cocoindex
54

65
# Define Qdrant connection constants
76
QDRANT_URL = "http://localhost:6334"
8-
QDRANT_COLLECTION = "cocoindex_text_embedding"
7+
QDRANT_COLLECTION = "TextEmbedding"
98

109

1110
@cocoindex.transform_flow()
@@ -55,15 +54,9 @@ def text_embedding_flow(
5554
text_embedding=chunk["embedding"],
5655
)
5756

58-
qdrant_conn = cocoindex.add_auth_entry(
59-
"Qdrant", cocoindex.storages.QdrantConnection(url=QDRANT_URL)
60-
)
6157
doc_embeddings.export(
6258
"doc_embeddings",
63-
cocoindex.storages.Qdrant(
64-
connection=qdrant_conn,
65-
collection_name=QDRANT_COLLECTION,
66-
),
59+
cocoindex.storages.Qdrant(collection_name=QDRANT_COLLECTION),
6760
primary_key_fields=["id"],
6861
)
6962

python/cocoindex/storages.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,16 @@ class Postgres(op.StorageSpec):
2020
class QdrantConnection:
2121
"""Connection spec for Qdrant."""
2222

23-
url: str
23+
grpc_url: str
2424
api_key: str | None = None
2525

2626

2727
@dataclass
2828
class Qdrant(op.StorageSpec):
2929
"""Storage powered by Qdrant - https://qdrant.tech/."""
3030

31-
connection: AuthEntryReference[QdrantConnection]
3231
collection_name: str
32+
connection: AuthEntryReference[QdrantConnection] | None = None
3333

3434

3535
@dataclass

src/ops/storages/qdrant.rs

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -24,13 +24,13 @@ const DEFAULT_URL: &str = "http://localhost:6334/";
2424

2525
#[derive(Debug, Deserialize, Clone)]
2626
pub struct ConnectionSpec {
27-
url: String,
27+
grpc_url: String,
2828
api_key: Option<String>,
2929
}
3030

3131
#[derive(Debug, Deserialize, Clone)]
3232
struct Spec {
33-
connection: spec::AuthEntryReference<ConnectionSpec>,
33+
connection: Option<spec::AuthEntryReference<ConnectionSpec>>,
3434
collection_name: String,
3535
}
3636

@@ -150,10 +150,14 @@ impl setup::ResourceSetupStatus for SetupStatus {
150150
}
151151

152152
impl SetupStatus {
153-
async fn apply(&self, collection_name: &String, qdrant_client: &Qdrant) -> Result<()> {
153+
async fn apply_delete(&self, collection_name: &String, qdrant_client: &Qdrant) -> Result<()> {
154154
if self.delete_collection {
155155
qdrant_client.delete_collection(collection_name).await?;
156156
}
157+
Ok(())
158+
}
159+
160+
async fn apply_create(&self, collection_name: &String, qdrant_client: &Qdrant) -> Result<()> {
157161
if let Some(add_collection) = &self.add_collection {
158162
let mut builder = CreateCollectionBuilder::new(collection_name);
159163
if !add_collection.vectors.is_empty() {
@@ -382,10 +386,9 @@ impl StorageFactoryBase for Factory {
382386
}
383387
}
384388

385-
let connection = Some(d.spec.connection);
386389
let export_context = Arc::new(ExportContext {
387390
qdrant_client: self
388-
.get_qdrant_client(&connection, &context.auth_registry)?,
391+
.get_qdrant_client(&d.spec.connection, &context.auth_registry)?,
389392
collection_name: d.spec.collection_name.clone(),
390393
fields_info,
391394
});
@@ -398,7 +401,7 @@ impl StorageFactoryBase for Factory {
398401
Ok(TypedExportDataCollectionBuildOutput {
399402
executors: executors.boxed(),
400403
setup_key: CollectionKey {
401-
connection,
404+
connection: d.spec.connection,
402405
collection_name: d.spec.collection_name,
403406
},
404407
desired_setup_state: SetupState {
@@ -489,12 +492,20 @@ impl StorageFactoryBase for Factory {
489492
setup_status: Vec<TypedResourceSetupChangeItem<'async_trait, Self>>,
490493
auth_registry: &Arc<AuthRegistry>,
491494
) -> Result<()> {
492-
for setup_change in setup_status.into_iter() {
495+
for setup_change in setup_status.iter() {
496+
let qdrant_client =
497+
self.get_qdrant_client(&setup_change.key.connection, auth_registry)?;
498+
setup_change
499+
.setup_status
500+
.apply_delete(&setup_change.key.collection_name, &qdrant_client)
501+
.await?;
502+
}
503+
for setup_change in setup_status.iter() {
493504
let qdrant_client =
494505
self.get_qdrant_client(&setup_change.key.connection, auth_registry)?;
495506
setup_change
496507
.setup_status
497-
.apply(&setup_change.key.collection_name, &qdrant_client)
508+
.apply_create(&setup_change.key.collection_name, &qdrant_client)
498509
.await?;
499510
}
500511
Ok(())
@@ -521,14 +532,14 @@ impl Factory {
521532
let spec = auth_entry.as_ref().map_or_else(
522533
|| {
523534
Ok(ConnectionSpec {
524-
url: DEFAULT_URL.to_string(),
535+
grpc_url: DEFAULT_URL.to_string(),
525536
api_key: None,
526537
})
527538
},
528539
|auth_entry| auth_registry.get(auth_entry),
529540
)?;
530541
let client = Arc::new(
531-
Qdrant::from_url(&spec.url)
542+
Qdrant::from_url(&spec.grpc_url)
532543
.api_key(spec.api_key)
533544
.skip_compatibility_check()
534545
.build()?,

0 commit comments

Comments
 (0)