Merge branch 'feat/custom-headers-builder' of github.com:qdrant/rust-client into feat/custom-headers-builder

Author: dancixx

tests/snippet_tests/mod.rs

@@ -21,6 +21,7 @@ mod test_delete_snapshot;
 mod test_delete_vectors;
 mod test_discover_batch_points;
 mod test_discover_points;
+mod test_external_api_keys;
 mod test_facets;
 mod test_get_collection;
 mod test_get_collection_aliases;
@@ -57,4 +58,4 @@ mod test_upsert_image;
 mod test_upsert_points;
 mod test_upsert_points_fallback_shard_key;
 mod test_upsert_points_insert_only;
-mod test_upsert_points_with_condition;
+mod test_upsert_points_with_condition;

tests/snippet_tests/test_external_api_keys.rs

@@ -0,0 +1,248 @@
+use std::collections::HashMap;
+
+use qdrant_client::qdrant::{
+    CreateCollectionBuilder, Distance, Document, PointStruct, Query, QueryPointsBuilder,
+    UpsertPointsBuilder, VectorParamsBuilder,
+};
+use qdrant_client::{Payload, Qdrant};
+use serde_json::json;
+
+const PROXY_URL: &str = "http://localhost:6334";
+const UPSERT_COLLECTION_NAME: &str = "test_external_api_keys_upsert";
+const QUERY_COLLECTION_NAME: &str = "test_external_api_keys_query";
+const DUAL_OPENAI_COLLECTION_NAME: &str = "test_external_api_keys_dual_openai";
+const DUAL_COHERE_COLLECTION_NAME: &str = "test_external_api_keys_dual_cohere";
+const OPENAI_MODEL: &str = "openai/text-embedding-3-small";
+const OPENAI_VECTOR_SIZE: u64 = 1536;
+const COHERE_MODEL: &str = "cohere/embed-english-v3.0";
+const COHERE_VECTOR_SIZE: u64 = 1024;
+
+fn create_client_with_external_keys(external_api_keys: HashMap<String, String>) -> Qdrant {
+    let mut builder = Qdrant::from_url(PROXY_URL)
+        .skip_compatibility_check()
+        .api_key("1234")
+        .timeout(30u64);
+    for (key, value) in external_api_keys {
+        builder = builder.header(key, value);
+    }
+    builder.build().expect("Failed to build client")
+}
+
+async fn setup_collection(client: &Qdrant, collection_name: &str, vector_size: u64) {
+    let _ = client.delete_collection(collection_name).await;
+
+    client
+        .create_collection(
+            CreateCollectionBuilder::new(collection_name)
+                .vectors_config(VectorParamsBuilder::new(vector_size, Distance::Cosine)),
+        )
+        .await
+        .expect("Failed to create collection");
+}
+
+fn cohere_document(text: impl Into<String>, input_type: &'static str) -> Document {
+    Document {
+        text: text.into(),
+        model: COHERE_MODEL.to_string(),
+        options: HashMap::from([("input_type".to_string(), input_type.into())]),
+    }
+}
+
+#[tokio::test]
+async fn test_upsert_with_external_api_keys() {
+    let Some(openai_api_key) = std::env::var("OPENAI_API_KEY").ok() else {
+        eprintln!("Skipping test_upsert_with_external_api_keys: OPENAI_API_KEY is not set");
+        return;
+    };
+    let collection_name = UPSERT_COLLECTION_NAME;
+    let client = create_client_with_external_keys(HashMap::from([(
+        "openai-api-key".to_string(),
+        openai_api_key,
+    )]));
+    setup_collection(&client, collection_name, OPENAI_VECTOR_SIZE).await;
+
+    let doc = Document::new("Qdrant is a vector search engine", OPENAI_MODEL);
+
+    let result = client
+        .upsert_points(
+            UpsertPointsBuilder::new(
+                collection_name,
+                vec![PointStruct::new(
+                    1,
+                    doc,
+                    Payload::try_from(json!({"source": "test"})).unwrap(),
+                )],
+            )
+            .wait(true),
+        )
+        .await;
+
+    assert!(
+        result.is_ok(),
+        "Upsert with external API keys failed: {result:?}"
+    );
+
+    let _ = client.delete_collection(collection_name).await;
+}
+
+#[tokio::test]
+async fn test_query_with_external_api_keys() {
+    let Some(openai_api_key) = std::env::var("OPENAI_API_KEY").ok() else {
+        eprintln!("Skipping test_query_with_external_api_keys: OPENAI_API_KEY is not set");
+        return;
+    };
+    let collection_name = QUERY_COLLECTION_NAME;
+    let client = create_client_with_external_keys(HashMap::from([(
+        "openai-api-key".to_string(),
+        openai_api_key,
+    )]));
+    setup_collection(&client, collection_name, OPENAI_VECTOR_SIZE).await;
+
+    // Upsert a point first
+    let doc = Document::new("Qdrant is a vector search engine", OPENAI_MODEL);
+    client
+        .upsert_points(
+            UpsertPointsBuilder::new(
+                collection_name,
+                vec![PointStruct::new(
+                    1,
+                    doc,
+                    Payload::try_from(json!({"source": "test"})).unwrap(),
+                )],
+            )
+            .wait(true),
+        )
+        .await
+        .expect("Upsert failed");
+
+    // Query with a document (server-side inference)
+    let query_doc = Document::new("vector database", OPENAI_MODEL);
+
+    let result = client
+        .query(
+            QueryPointsBuilder::new(collection_name)
+                .query(Query::new_nearest(query_doc))
+                .limit(1)
+                .with_payload(true),
+        )
+        .await;
+
+    assert!(
+        result.is_ok(),
+        "Query with external API keys failed: {result:?}"
+    );
+
+    let response = result.unwrap();
+    assert_eq!(response.result.len(), 1);
+    assert!(response.result[0].payload.contains_key("source"));
+
+    let _ = client.delete_collection(collection_name).await;
+}
+
+#[tokio::test]
+async fn test_query_with_two_external_api_providers() {
+    let Some(openai_api_key) = std::env::var("OPENAI_API_KEY").ok() else {
+        eprintln!("Skipping test_query_with_two_external_api_providers: OPENAI_API_KEY is not set");
+        return;
+    };
+    let Some(cohere_api_key) = std::env::var("COHERE_API_KEY").ok() else {
+        eprintln!("Skipping test_query_with_two_external_api_providers: COHERE_API_KEY is not set");
+        return;
+    };
+
+    let client = create_client_with_external_keys(HashMap::from([
+        ("openai-api-key".to_string(), openai_api_key),
+        ("cohere-api-key".to_string(), cohere_api_key),
+    ]));
+
+    setup_collection(&client, DUAL_OPENAI_COLLECTION_NAME, OPENAI_VECTOR_SIZE).await;
+    setup_collection(&client, DUAL_COHERE_COLLECTION_NAME, COHERE_VECTOR_SIZE).await;
+
+    let openai_doc = Document::new("OpenAI provider document", OPENAI_MODEL);
+    let cohere_doc = cohere_document("Cohere provider document", "search_document");
+
+    let openai_upsert = client
+        .upsert_points(
+            UpsertPointsBuilder::new(
+                DUAL_OPENAI_COLLECTION_NAME,
+                vec![PointStruct::new(
+                    1,
+                    openai_doc,
+                    Payload::try_from(json!({"provider": "openai"})).unwrap(),
+                )],
+            )
+            .wait(true),
+        )
+        .await;
+    assert!(
+        openai_upsert.is_ok(),
+        "OpenAI upsert with external API keys failed: {openai_upsert:?}"
+    );
+
+    let cohere_upsert = client
+        .upsert_points(
+            UpsertPointsBuilder::new(
+                DUAL_COHERE_COLLECTION_NAME,
+                vec![PointStruct::new(
+                    1,
+                    cohere_doc,
+                    Payload::try_from(json!({"provider": "cohere"})).unwrap(),
+                )],
+            )
+            .wait(true),
+        )
+        .await;
+    assert!(
+        cohere_upsert.is_ok(),
+        "Cohere upsert with external API keys failed: {cohere_upsert:?}"
+    );
+
+    let openai_query = client
+        .query(
+            QueryPointsBuilder::new(DUAL_OPENAI_COLLECTION_NAME)
+                .query(Query::new_nearest(Document::new(
+                    "OpenAI provider query",
+                    OPENAI_MODEL,
+                )))
+                .limit(1)
+                .with_payload(true),
+        )
+        .await;
+    assert!(
+        openai_query.is_ok(),
+        "OpenAI query with external API keys failed: {openai_query:?}"
+    );
+
+    let cohere_query = client
+        .query(
+            QueryPointsBuilder::new(DUAL_COHERE_COLLECTION_NAME)
+                .query(Query::new_nearest(cohere_document(
+                    "Cohere provider query",
+                    "search_query",
+                )))
+                .limit(1)
+                .with_payload(true),
+        )
+        .await;
+    assert!(
+        cohere_query.is_ok(),
+        "Cohere query with external API keys failed: {cohere_query:?}"
+    );
+
+    let openai_response = openai_query.unwrap();
+    assert_eq!(openai_response.result.len(), 1);
+    assert_eq!(
+        openai_response.result[0].payload["provider"],
+        "openai".into()
+    );
+
+    let cohere_response = cohere_query.unwrap();
+    assert_eq!(cohere_response.result.len(), 1);
+    assert_eq!(
+        cohere_response.result[0].payload["provider"],
+        "cohere".into()
+    );
+
+    let _ = client.delete_collection(DUAL_OPENAI_COLLECTION_NAME).await;
+    let _ = client.delete_collection(DUAL_COHERE_COLLECTION_NAME).await;
+}