diff --git a/tests/snippet_tests/mod.rs b/tests/snippet_tests/mod.rs index decfd24..b33ed5e 100644 --- a/tests/snippet_tests/mod.rs +++ b/tests/snippet_tests/mod.rs @@ -20,6 +20,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; @@ -56,4 +57,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; \ No newline at end of file +mod test_upsert_points_with_condition; diff --git a/tests/snippet_tests/test_external_api_keys.rs b/tests/snippet_tests/test_external_api_keys.rs new file mode 100644 index 0000000..ebb4473 --- /dev/null +++ b/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) -> 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, 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; +}