From e6d95374a2553d3aff01d032fad40a9443908ad9 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Tue, 4 Jun 2024 16:15:14 -0400 Subject: [PATCH 01/20] added source tag cleaning --- Cargo.lock | 1 + pinecone_sdk/Cargo.toml | 3 ++- pinecone_sdk/src/lib.rs | 1 + pinecone_sdk/src/utils/mod.rs | 1 + pinecone_sdk/src/utils/user_agent.rs | 30 ++++++++++++++++++++++++++++ 5 files changed, 35 insertions(+), 1 deletion(-) create mode 100644 pinecone_sdk/src/utils/mod.rs create mode 100644 pinecone_sdk/src/utils/user_agent.rs diff --git a/Cargo.lock b/Cargo.lock index ac01c01..b10f969 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -680,6 +680,7 @@ version = "0.1.0" dependencies = [ "mockito", "openapi", + "regex", "tokio", ] diff --git a/pinecone_sdk/Cargo.toml b/pinecone_sdk/Cargo.toml index d4ce97d..08f6aee 100644 --- a/pinecone_sdk/Cargo.toml +++ b/pinecone_sdk/Cargo.toml @@ -8,4 +8,5 @@ edition = "2021" [dependencies] openapi = { path = "../openapi" } tokio = { version = "1", features = ["full"] } -mockito = "0.30" \ No newline at end of file +mockito = "0.30" +regex = "1.10.4" diff --git a/pinecone_sdk/src/lib.rs b/pinecone_sdk/src/lib.rs index 52cb989..68cf8b9 100644 --- a/pinecone_sdk/src/lib.rs +++ b/pinecone_sdk/src/lib.rs @@ -1,3 +1,4 @@ pub mod config; pub mod control; pub mod pinecone; +pub mod utils; diff --git a/pinecone_sdk/src/utils/mod.rs b/pinecone_sdk/src/utils/mod.rs new file mode 100644 index 0000000..0e90bb1 --- /dev/null +++ b/pinecone_sdk/src/utils/mod.rs @@ -0,0 +1 @@ +mod user_agent; \ No newline at end of file diff --git a/pinecone_sdk/src/utils/user_agent.rs b/pinecone_sdk/src/utils/user_agent.rs new file mode 100644 index 0000000..9c01ea7 --- /dev/null +++ b/pinecone_sdk/src/utils/user_agent.rs @@ -0,0 +1,30 @@ +use regex::Regex; + +// Normalizes the source tag +fn build_source_tag(source_tag: String) -> String { + // 1. Lowercase + // 2. Limit charset to [a-z0-9_ ] + // 3. Trim left/right empty space + // 4. Condence multiple spaces to one, and replace with underscore + + let re = Regex::new(r"[^a-z0-9_ ]").unwrap(); + let lowercase_tag = source_tag.to_lowercase(); + let tag = re.replace_all(&lowercase_tag, ""); + return tag.trim() + .split(' ') + .filter(|s| !s.is_empty()) + .collect::>() + .join("_"); +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio; + + #[tokio::test] + async fn test_build_source_tag() { + let source_tag = " Hello World!! ".to_string(); + assert_eq!(build_source_tag(source_tag), "hello_world"); + } +} \ No newline at end of file From 7d4ea0e0f7a152a7aa3bd60c90da40724a737ba7 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Wed, 5 Jun 2024 10:00:16 -0400 Subject: [PATCH 02/20] added get_user_string function --- pinecone_sdk/src/config.rs | 2 ++ pinecone_sdk/src/pinecone.rs | 5 +++- pinecone_sdk/src/utils/mod.rs | 2 +- pinecone_sdk/src/utils/user_agent.rs | 34 ++++++++++++++++++++++++++-- 4 files changed, 39 insertions(+), 4 deletions(-) diff --git a/pinecone_sdk/src/config.rs b/pinecone_sdk/src/config.rs index 500d9e5..0c80ab3 100644 --- a/pinecone_sdk/src/config.rs +++ b/pinecone_sdk/src/config.rs @@ -2,6 +2,7 @@ pub struct Config { pub api_key: String, pub controller_url: String, + pub source_tag: Option, } impl Config { @@ -9,6 +10,7 @@ impl Config { Config { api_key, controller_url: "https://api.pinecone.io".to_string(), + source_tag: None, } } } diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index 2d91522..e42dbbd 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -1,4 +1,5 @@ use crate::config::Config; +use crate::utils::user_agent::get_user_agent; use openapi::apis::configuration::ApiKey; use openapi::apis::configuration::Configuration; @@ -12,9 +13,11 @@ impl Pinecone { pub fn new(api_key: String, control_plane_host: Option) -> Self { let config = Config::new(api_key.clone()); + let user_agent = get_user_agent(&config); + let openapi_config = Configuration { base_path: control_plane_host.unwrap_or("https://api.pinecone.io".to_string()), - user_agent: Some("pinecone-rust-client".to_string()), + user_agent: Some(user_agent), api_key: Some(ApiKey { prefix: None, key: api_key, diff --git a/pinecone_sdk/src/utils/mod.rs b/pinecone_sdk/src/utils/mod.rs index 0e90bb1..c497f92 100644 --- a/pinecone_sdk/src/utils/mod.rs +++ b/pinecone_sdk/src/utils/mod.rs @@ -1 +1 @@ -mod user_agent; \ No newline at end of file +pub mod user_agent; \ No newline at end of file diff --git a/pinecone_sdk/src/utils/user_agent.rs b/pinecone_sdk/src/utils/user_agent.rs index 9c01ea7..5726c42 100644 --- a/pinecone_sdk/src/utils/user_agent.rs +++ b/pinecone_sdk/src/utils/user_agent.rs @@ -1,7 +1,8 @@ use regex::Regex; +use crate::config::Config; // Normalizes the source tag -fn build_source_tag(source_tag: String) -> String { +fn build_source_tag(source_tag: &String) -> String { // 1. Lowercase // 2. Limit charset to [a-z0-9_ ] // 3. Trim left/right empty space @@ -17,6 +18,15 @@ fn build_source_tag(source_tag: String) -> String { .join("_"); } +// Gets user agent string +pub fn get_user_agent(config: &Config) -> String { + let mut user_agent = format!("lang=rust/{}", "0.1.0"); + if let Some(source_tag) = &config.source_tag { + user_agent.push_str(&format!("; source_tag={}", build_source_tag(source_tag))); + } + return user_agent; +} + #[cfg(test)] mod tests { use super::*; @@ -25,6 +35,26 @@ mod tests { #[tokio::test] async fn test_build_source_tag() { let source_tag = " Hello World!! ".to_string(); - assert_eq!(build_source_tag(source_tag), "hello_world"); + assert_eq!(build_source_tag(&source_tag), "hello_world"); + } + + #[tokio::test] + async fn test_no_source_tag() { + let config = Config { + api_key: "api".to_string(), + controller_url: "https://api.pinecone.io".to_string(), + source_tag: None, + }; + assert_eq!(get_user_agent(&config), "lang=rust/0.1.0"); + } + + #[tokio::test] + async fn test_with_source_tag() { + let config = Config { + api_key: "api".to_string(), + controller_url: "https://api.pinecone.io".to_string(), + source_tag: Some("tag".to_string()), + }; + assert_eq!(get_user_agent(&config), "lang=rust/0.1.0; source_tag=tag"); } } \ No newline at end of file From bbf12b9733f6bc7432bdbdfb8625b2ec174ee0c8 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Thu, 6 Jun 2024 13:59:56 -0400 Subject: [PATCH 03/20] added create_serverless_index function --- pinecone_sdk/src/control/create_index.rs | 145 ++++++++++++++++++ pinecone_sdk/src/control/mod.rs | 3 +- pinecone_sdk/src/lib.rs | 1 + .../src/models/create_index_request_params.rs | 20 +++ pinecone_sdk/src/models/mod.rs | 1 + 5 files changed, 169 insertions(+), 1 deletion(-) create mode 100644 pinecone_sdk/src/control/create_index.rs create mode 100644 pinecone_sdk/src/models/create_index_request_params.rs create mode 100644 pinecone_sdk/src/models/mod.rs diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs new file mode 100644 index 0000000..4289add --- /dev/null +++ b/pinecone_sdk/src/control/create_index.rs @@ -0,0 +1,145 @@ +use core::panic; +use openapi::apis::manage_indexes_api::CreateIndexError; +use openapi::apis::Error; +use openapi::models::{CreateIndexRequest, CreateIndexRequestSpec, IndexModel, ServerlessSpec}; +use openapi::models::create_index_request::Metric; +use openapi::models::serverless_spec::Cloud; +use crate::pinecone::Pinecone; +use crate::models::create_index_request_params::CreateServerlessIndexRequest; + +impl Pinecone { + pub async fn create_serverless_index(&self, params: CreateServerlessIndexRequest) -> Result> { + let create_index_request = self.create_serverless_index_req(params); + let response = openapi::apis::manage_indexes_api::create_index( + &self.openapi_config(), + create_index_request + ).await?; + Ok(response) + } + + pub fn create_serverless_index_req(&self, params: CreateServerlessIndexRequest) -> CreateIndexRequest { + // clean metric string + let metric_enum = match ¶ms.metric { + Some(metric) => match metric.as_str() { + "cosine" => Some(Metric::Cosine), + "euclidean" => Some(Metric::Euclidean), + "dotproduct" => Some(Metric::Dotproduct), + _ => panic!("Invalid metric"), + }, + None => None, + }; + + // clean cloud string + let cloud_enum = match ¶ms.cloud { + Some(cloud) => match cloud.as_str() { + "gcp" => Cloud::Gcp, + "aws" => Cloud::Aws, + "azure" => Cloud::Azure, + _ => panic!("Invalid cloud type"), + }, + None => Cloud::default(), + }; + + // create request specs + let create_index_request_spec = CreateIndexRequestSpec { + serverless: Some(Box::new(ServerlessSpec { + cloud: cloud_enum, + region: params.region, + })), + pod: None, + }; + + let create_index_request = CreateIndexRequest { + name: params.name, + dimension: params.dimension, + metric: metric_enum, + spec: Some(Box::new(create_index_request_spec)), + }; + + return create_index_request; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use core::panic; + use mockito::mock; + use openapi::models::serverless_spec; + use tokio; + + #[tokio::test] + async fn test_create_serverless_index_req() { + let pinecone = Pinecone::new("api_key".to_string(), Some("controller_url".to_string())); + let params = CreateServerlessIndexRequest { + name: "index_name".to_string(), + dimension: 10, + metric: Some("cosine".to_string()), + cloud: Some("gcp".to_string()), + region: "us-east-1".to_string(), + }; + + let create_index_request = pinecone.create_serverless_index_req(params); + assert_eq!(create_index_request.name, "index_name"); + assert_eq!(create_index_request.dimension, 10); + assert_eq!(create_index_request.metric, Some(Metric::Cosine)); + + let spec = create_index_request.spec.unwrap(); + let serverless_spec = spec.serverless.unwrap(); + assert_eq!(serverless_spec.cloud, serverless_spec::Cloud::Gcp); + assert_eq!(serverless_spec.region, "us-east-1"); + } + + #[tokio::test] + async fn test_create_serverless_index() { + // Create a mock server + let _m = mock("POST", "/indexes") + .with_status(201) + .with_header("content-type", "application/json") + .with_body( + r#" + { + "name": "index_name", + "dimension": 10, + "metric": "euclidean", + "host": "host1", + "spec": { + "serverless": { + "cloud": "aws", + "region": "us-east-1" + } + }, + "status": { + "ready": true, + "state": "Initializing" + } + } + "#, + ) + .create(); + + let pinecone = Pinecone::new("api_key".to_string(), Some(mockito::server_url())); + let params = CreateServerlessIndexRequest { + name: "index_name".to_string(), + dimension: 10, + metric: Some("euclidean".to_string()), + cloud: Some("aws".to_string()), + region: "us-east-1".to_string(), + }; + + let result = pinecone.create_serverless_index(params).await; + + match result { + Ok(index) => { + assert_eq!(index.name, "index_name"); + assert_eq!(index.dimension, 10); + assert_eq!(index.metric, openapi::models::index_model::Metric::Euclidean); + let spec = *index.spec; + let serverless_spec = spec.serverless.unwrap(); + assert_eq!(serverless_spec.cloud, openapi::models::serverless_spec::Cloud::Aws); + assert_eq!(serverless_spec.region, "us-east-1"); + }, + Err(e) => panic!("{}", e), + } + } +} \ No newline at end of file diff --git a/pinecone_sdk/src/control/mod.rs b/pinecone_sdk/src/control/mod.rs index bca88fd..9821562 100644 --- a/pinecone_sdk/src/control/mod.rs +++ b/pinecone_sdk/src/control/mod.rs @@ -1,3 +1,4 @@ -mod list_indexes; +pub mod create_index; +pub mod list_indexes; pub use list_indexes::*; diff --git a/pinecone_sdk/src/lib.rs b/pinecone_sdk/src/lib.rs index 68cf8b9..f26d99b 100644 --- a/pinecone_sdk/src/lib.rs +++ b/pinecone_sdk/src/lib.rs @@ -2,3 +2,4 @@ pub mod config; pub mod control; pub mod pinecone; pub mod utils; +pub mod models; \ No newline at end of file diff --git a/pinecone_sdk/src/models/create_index_request_params.rs b/pinecone_sdk/src/models/create_index_request_params.rs new file mode 100644 index 0000000..02dede4 --- /dev/null +++ b/pinecone_sdk/src/models/create_index_request_params.rs @@ -0,0 +1,20 @@ + +pub struct CreateServerlessIndexRequest { + pub name: String, + pub dimension: i32, + pub metric: Option, + pub cloud: Option, + pub region: String, +} + +impl CreateServerlessIndexRequest { + pub fn new(name: &str, dimension: i32, metric: Option<&str>, cloud: Option<&str>, region: &str) -> CreateServerlessIndexRequest { + CreateServerlessIndexRequest { + name: name.to_string(), + dimension, + metric: if let Some(metric) = metric { Some(metric.to_string()) } else { None }, + cloud: if let Some(cloud) = cloud { Some(cloud.to_string()) } else { None }, + region: region.to_string(), + } + } +} \ No newline at end of file diff --git a/pinecone_sdk/src/models/mod.rs b/pinecone_sdk/src/models/mod.rs new file mode 100644 index 0000000..b1fe893 --- /dev/null +++ b/pinecone_sdk/src/models/mod.rs @@ -0,0 +1 @@ +pub mod create_index_request_params; \ No newline at end of file From 0bc7db9599e7084ece52a77f433a71e9cd2241d0 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Thu, 6 Jun 2024 15:19:45 -0400 Subject: [PATCH 04/20] added some notes for future --- pinecone_sdk/src/control/create_index.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs index 4289add..c76b8d3 100644 --- a/pinecone_sdk/src/control/create_index.rs +++ b/pinecone_sdk/src/control/create_index.rs @@ -24,7 +24,7 @@ impl Pinecone { "cosine" => Some(Metric::Cosine), "euclidean" => Some(Metric::Euclidean), "dotproduct" => Some(Metric::Dotproduct), - _ => panic!("Invalid metric"), + _ => panic!("Invalid metric"), // TODO: handle error better }, None => None, }; @@ -35,7 +35,7 @@ impl Pinecone { "gcp" => Cloud::Gcp, "aws" => Cloud::Aws, "azure" => Cloud::Azure, - _ => panic!("Invalid cloud type"), + _ => panic!("Invalid cloud type"), // TODO: handle error better }, None => Cloud::default(), }; From 35fbed4f9b9e4780b1dc47fa23c9f54319003fd4 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Fri, 7 Jun 2024 09:59:16 -0400 Subject: [PATCH 05/20] fixed more merge conflicts and tests --- Cargo.lock | 3 --- pinecone_sdk/Cargo.toml | 5 ----- pinecone_sdk/src/control/create_index.rs | 10 ++++++---- pinecone_sdk/src/pinecone.rs | 2 ++ 4 files changed, 8 insertions(+), 12 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 892e5aa..1c2289e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -725,12 +725,9 @@ dependencies = [ "mockito", "openapi", "regex", -<<<<<<< HEAD -======= "serde_json", "serial_test", "snafu", ->>>>>>> origin/main "tokio", ] diff --git a/pinecone_sdk/Cargo.toml b/pinecone_sdk/Cargo.toml index b6a6fa4..3bdebc4 100644 --- a/pinecone_sdk/Cargo.toml +++ b/pinecone_sdk/Cargo.toml @@ -8,10 +8,6 @@ edition = "2021" [dependencies] openapi = { path = "../openapi" } tokio = { version = "1", features = ["full"] } -<<<<<<< HEAD -mockito = "0.30" -regex = "1.10.4" -======= regex = "1.10.4" serde_json = "1.0.117" snafu = "0.8.3" @@ -19,4 +15,3 @@ snafu = "0.8.3" [dev-dependencies] mockito = "0.30" serial_test = "3.1.1" ->>>>>>> origin/main diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs index c76b8d3..aba595f 100644 --- a/pinecone_sdk/src/control/create_index.rs +++ b/pinecone_sdk/src/control/create_index.rs @@ -66,11 +66,13 @@ mod tests { use core::panic; use mockito::mock; use openapi::models::serverless_spec; + use serial_test::serial; use tokio; #[tokio::test] + #[serial] async fn test_create_serverless_index_req() { - let pinecone = Pinecone::new("api_key".to_string(), Some("controller_url".to_string())); + let pinecone = Pinecone::new(Some("api_key".to_string()), Some("controller_url".to_string()), None, None); let params = CreateServerlessIndexRequest { name: "index_name".to_string(), dimension: 10, @@ -79,7 +81,7 @@ mod tests { region: "us-east-1".to_string(), }; - let create_index_request = pinecone.create_serverless_index_req(params); + let create_index_request = pinecone.expect("REASON").create_serverless_index_req(params); assert_eq!(create_index_request.name, "index_name"); assert_eq!(create_index_request.dimension, 10); assert_eq!(create_index_request.metric, Some(Metric::Cosine)); @@ -118,7 +120,7 @@ mod tests { ) .create(); - let pinecone = Pinecone::new("api_key".to_string(), Some(mockito::server_url())); + let pinecone = Pinecone::new(Some("api_key".to_string()), Some(mockito::server_url()), None, None); let params = CreateServerlessIndexRequest { name: "index_name".to_string(), dimension: 10, @@ -127,7 +129,7 @@ mod tests { region: "us-east-1".to_string(), }; - let result = pinecone.create_serverless_index(params).await; + let result = pinecone.expect("REASON").create_serverless_index(params).await; match result { Ok(index) => { diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index 23e4751..c282f7f 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -278,6 +278,7 @@ mod tests { pinecone.err().unwrap(), PineconeError::InvalidHeadersError { .. } )); + remove_env_var("PINECONE_ADDITIONAL_HEADERS"); } #[tokio::test] @@ -345,5 +346,6 @@ mod tests { pinecone.as_ref().unwrap().config.additional_headers, mock_arg_headers.clone() ); + } } From d0689b02b7d8c55390682c96d5aafda3f3511e0c Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Fri, 7 Jun 2024 10:05:04 -0400 Subject: [PATCH 06/20] started error checking --- pinecone_sdk/src/pinecone.rs | 4 +--- pinecone_sdk/src/utils/errors.rs | 6 ++++++ 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index c282f7f..e52c044 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -54,9 +54,7 @@ impl Pinecone { additional_headers, source_tag, }; - - let user_agent = get_user_agent(&config); - + let user_agent = get_user_agent(&config); let openapi_config = Configuration { diff --git a/pinecone_sdk/src/utils/errors.rs b/pinecone_sdk/src/utils/errors.rs index e97c306..dc1c7e2 100644 --- a/pinecone_sdk/src/utils/errors.rs +++ b/pinecone_sdk/src/utils/errors.rs @@ -7,4 +7,10 @@ pub enum PineconeError { #[snafu(display("Failed to parse headers: {}", json_error))] InvalidHeadersError { json_error: serde_json::Error }, + + #[snafu(display("Invalid region."))] + InvalidRegionError, + + #[snafu(display("Invalid cloud."))] + InvalidCloudError, } From c48956bb47665883179353d87bfd77a97193cc6b Mon Sep 17 00:00:00 2001 From: Emily Yu Date: Fri, 7 Jun 2024 11:49:51 -0400 Subject: [PATCH 07/20] Add error handling to create_index --- pinecone_sdk/src/control/create_index.rs | 106 ++++++++++++++--------- pinecone_sdk/src/control/list_indexes.rs | 3 - pinecone_sdk/src/control/mod.rs | 2 - pinecone_sdk/src/utils/errors.rs | 15 +++- 4 files changed, 79 insertions(+), 47 deletions(-) diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs index aba595f..f8b34d5 100644 --- a/pinecone_sdk/src/control/create_index.rs +++ b/pinecone_sdk/src/control/create_index.rs @@ -1,44 +1,55 @@ -use core::panic; -use openapi::apis::manage_indexes_api::CreateIndexError; -use openapi::apis::Error; -use openapi::models::{CreateIndexRequest, CreateIndexRequestSpec, IndexModel, ServerlessSpec}; +use crate::models::create_index_request_params::CreateServerlessIndexRequest; +use crate::pinecone::Pinecone; +use crate::utils::errors::PineconeError; use openapi::models::create_index_request::Metric; use openapi::models::serverless_spec::Cloud; -use crate::pinecone::Pinecone; -use crate::models::create_index_request_params::CreateServerlessIndexRequest; +use openapi::models::{CreateIndexRequest, CreateIndexRequestSpec, IndexModel, ServerlessSpec}; impl Pinecone { - pub async fn create_serverless_index(&self, params: CreateServerlessIndexRequest) -> Result> { + pub async fn create_serverless_index( + &self, + params: CreateServerlessIndexRequest, + ) -> Result { let create_index_request = self.create_serverless_index_req(params); - let response = openapi::apis::manage_indexes_api::create_index( + match openapi::apis::manage_indexes_api::create_index( &self.openapi_config(), - create_index_request - ).await?; - Ok(response) + create_index_request?, + ) + .await + { + Ok(index) => Ok(index), + Err(e) => Err(PineconeError::CreateIndexError { openapi_error: e }), + } } - pub fn create_serverless_index_req(&self, params: CreateServerlessIndexRequest) -> CreateIndexRequest { - // clean metric string + pub fn create_serverless_index_req( + &self, + params: CreateServerlessIndexRequest, + ) -> Result { let metric_enum = match ¶ms.metric { Some(metric) => match metric.as_str() { - "cosine" => Some(Metric::Cosine), - "euclidean" => Some(Metric::Euclidean), - "dotproduct" => Some(Metric::Dotproduct), - _ => panic!("Invalid metric"), // TODO: handle error better + "cosine" => Ok(Some(Metric::Cosine)), + "euclidean" => Ok(Some(Metric::Euclidean)), + "dotproduct" => Ok(Some(Metric::Dotproduct)), + _ => Err(PineconeError::InvalidMetricError { + metric: metric.clone(), + }), }, - None => None, - }; + None => Ok(Some(Metric::Cosine)), + }?; // clean cloud string let cloud_enum = match ¶ms.cloud { Some(cloud) => match cloud.as_str() { - "gcp" => Cloud::Gcp, - "aws" => Cloud::Aws, - "azure" => Cloud::Azure, - _ => panic!("Invalid cloud type"), // TODO: handle error better + "gcp" => Ok(Cloud::Gcp), + "aws" => Ok(Cloud::Aws), + "azure" => Ok(Cloud::Azure), + _ => Err(PineconeError::InvalidCloudError { + cloud: cloud.clone(), + }), }, - None => Cloud::default(), - }; + None => Ok(Cloud::default()), + }?; // create request specs let create_index_request_spec = CreateIndexRequestSpec { @@ -48,15 +59,13 @@ impl Pinecone { })), pod: None, }; - - let create_index_request = CreateIndexRequest { + + Ok(CreateIndexRequest { name: params.name, dimension: params.dimension, metric: metric_enum, spec: Some(Box::new(create_index_request_spec)), - }; - - return create_index_request; + }) } } @@ -72,7 +81,12 @@ mod tests { #[tokio::test] #[serial] async fn test_create_serverless_index_req() { - let pinecone = Pinecone::new(Some("api_key".to_string()), Some("controller_url".to_string()), None, None); + let pinecone = Pinecone::new( + Some("api_key".to_string()), + Some("controller_url".to_string()), + None, + None, + ); let params = CreateServerlessIndexRequest { name: "index_name".to_string(), dimension: 10, @@ -81,7 +95,10 @@ mod tests { region: "us-east-1".to_string(), }; - let create_index_request = pinecone.expect("REASON").create_serverless_index_req(params); + let create_index_request = pinecone.unwrap().create_serverless_index_req(params); + assert!(create_index_request.is_ok()); + + let create_index_request = create_index_request.unwrap(); assert_eq!(create_index_request.name, "index_name"); assert_eq!(create_index_request.dimension, 10); assert_eq!(create_index_request.metric, Some(Metric::Cosine)); @@ -120,7 +137,12 @@ mod tests { ) .create(); - let pinecone = Pinecone::new(Some("api_key".to_string()), Some(mockito::server_url()), None, None); + let pinecone = Pinecone::new( + Some("api_key".to_string()), + Some(mockito::server_url()), + None, + None, + ); let params = CreateServerlessIndexRequest { name: "index_name".to_string(), dimension: 10, @@ -129,19 +151,25 @@ mod tests { region: "us-east-1".to_string(), }; - let result = pinecone.expect("REASON").create_serverless_index(params).await; - + let result = pinecone.unwrap().create_serverless_index(params).await; + match result { Ok(index) => { assert_eq!(index.name, "index_name"); assert_eq!(index.dimension, 10); - assert_eq!(index.metric, openapi::models::index_model::Metric::Euclidean); + assert_eq!( + index.metric, + openapi::models::index_model::Metric::Euclidean + ); let spec = *index.spec; let serverless_spec = spec.serverless.unwrap(); - assert_eq!(serverless_spec.cloud, openapi::models::serverless_spec::Cloud::Aws); + assert_eq!( + serverless_spec.cloud, + openapi::models::serverless_spec::Cloud::Aws + ); assert_eq!(serverless_spec.region, "us-east-1"); - }, + } Err(e) => panic!("{}", e), } } -} \ No newline at end of file +} diff --git a/pinecone_sdk/src/control/list_indexes.rs b/pinecone_sdk/src/control/list_indexes.rs index 6b9c1ab..4cea7ba 100644 --- a/pinecone_sdk/src/control/list_indexes.rs +++ b/pinecone_sdk/src/control/list_indexes.rs @@ -15,11 +15,8 @@ impl Pinecone { #[cfg(test)] mod tests { use super::*; - use crate::config::Config; use crate::control::list_indexes::models::index_model::Metric; use mockito::mock; - use openapi::apis::configuration::ApiKey; - use openapi::apis::configuration::Configuration; use openapi::models::IndexList; use openapi::models::IndexModel; use tokio; diff --git a/pinecone_sdk/src/control/mod.rs b/pinecone_sdk/src/control/mod.rs index 9821562..5b4c3c3 100644 --- a/pinecone_sdk/src/control/mod.rs +++ b/pinecone_sdk/src/control/mod.rs @@ -1,4 +1,2 @@ pub mod create_index; pub mod list_indexes; - -pub use list_indexes::*; diff --git a/pinecone_sdk/src/utils/errors.rs b/pinecone_sdk/src/utils/errors.rs index dc1c7e2..639b3dd 100644 --- a/pinecone_sdk/src/utils/errors.rs +++ b/pinecone_sdk/src/utils/errors.rs @@ -1,3 +1,4 @@ +use openapi::apis::{manage_indexes_api::CreateIndexError, Error as OpenAPIError}; use snafu::prelude::*; #[derive(Debug, Snafu)] @@ -5,12 +6,20 @@ pub enum PineconeError { #[snafu(display("API key missing."))] APIKeyMissingError, + #[snafu(display("API key missing."))] + CreateIndexError { + openapi_error: OpenAPIError, + }, + #[snafu(display("Failed to parse headers: {}", json_error))] InvalidHeadersError { json_error: serde_json::Error }, + #[snafu(display("Invalid cloud '{}'.", cloud))] + InvalidCloudError { cloud: String }, + + #[snafu(display("Invalid metric '{}'.", metric))] + InvalidMetricError { metric: String }, + #[snafu(display("Invalid region."))] InvalidRegionError, - - #[snafu(display("Invalid cloud."))] - InvalidCloudError, } From 0b177dbd4d9070db791730d30a54ce7cf1f8d530 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Fri, 7 Jun 2024 15:01:26 -0400 Subject: [PATCH 08/20] combined create_index for serverless and pods into one function, created struct/enum for spec, metric, cloud --- pinecone_sdk/src/control/create_index.rs | 191 +++++++++++------- .../src/models/create_index_request_params.rs | 54 +++-- 2 files changed, 161 insertions(+), 84 deletions(-) diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs index f8b34d5..e0ab7a6 100644 --- a/pinecone_sdk/src/control/create_index.rs +++ b/pinecone_sdk/src/control/create_index.rs @@ -1,71 +1,92 @@ -use crate::models::create_index_request_params::CreateServerlessIndexRequest; +use crate::models::create_index_request_params::{CreateIndexParams, Spec, Metric, Cloud}; use crate::pinecone::Pinecone; use crate::utils::errors::PineconeError; -use openapi::models::create_index_request::Metric; -use openapi::models::serverless_spec::Cloud; +use openapi::models::create_index_request; +use openapi::models::serverless_spec; use openapi::models::{CreateIndexRequest, CreateIndexRequestSpec, IndexModel, ServerlessSpec}; impl Pinecone { - pub async fn create_serverless_index( + // Creates a new index based on the provided parameters + pub async fn create_index( &self, - params: CreateServerlessIndexRequest, + params: CreateIndexParams ) -> Result { - let create_index_request = self.create_serverless_index_req(params); - match openapi::apis::manage_indexes_api::create_index( - &self.openapi_config(), - create_index_request?, - ) - .await - { - Ok(index) => Ok(index), - Err(e) => Err(PineconeError::CreateIndexError { openapi_error: e }), + + // Check if creating serverless or pod-based index and call respective builder function + match params.spec { + Spec::Serverless { cloud, region } => { + self.create_serverless_index( + params.name, + params.dimension, + params.metric, + cloud, + region + ).await + } + Spec::Pod { + environment: _, + replicas: _, + shards: _, + pod_type: _, + pods: _, + metadata_config: _, + source_collection: _, + } => { + // eventually change this to be pod index + self.create_serverless_index(params.name, params.dimension, params.metric, Cloud::Aws, "".to_string()).await + } } } - pub fn create_serverless_index_req( + // Creates serverless index + async fn create_serverless_index( &self, - params: CreateServerlessIndexRequest, - ) -> Result { - let metric_enum = match ¶ms.metric { - Some(metric) => match metric.as_str() { - "cosine" => Ok(Some(Metric::Cosine)), - "euclidean" => Ok(Some(Metric::Euclidean)), - "dotproduct" => Ok(Some(Metric::Dotproduct)), - _ => Err(PineconeError::InvalidMetricError { - metric: metric.clone(), - }), - }, - None => Ok(Some(Metric::Cosine)), - }?; - - // clean cloud string - let cloud_enum = match ¶ms.cloud { - Some(cloud) => match cloud.as_str() { - "gcp" => Ok(Cloud::Gcp), - "aws" => Ok(Cloud::Aws), - "azure" => Ok(Cloud::Azure), - _ => Err(PineconeError::InvalidCloudError { - cloud: cloud.clone(), - }), - }, - None => Ok(Cloud::default()), - }?; + name: String, + dimension: u32, + metric: Metric, + cloud: Cloud, + region: String + ) -> Result { + + // clean metric enum + let metric_enum = match metric { + Metric::Cosine => Some(create_index_request::Metric::Cosine), + Metric::Dotproduct => Some(create_index_request::Metric::Dotproduct), + Metric::Euclidean => Some(create_index_request::Metric::Euclidean), + }; + + // clean cloud enum + let cloud_enum = match cloud { + Cloud::Aws => serverless_spec::Cloud::Aws, + Cloud::Gcp => serverless_spec::Cloud::Gcp, + Cloud::Azure => serverless_spec::Cloud::Azure, + }; // create request specs let create_index_request_spec = CreateIndexRequestSpec { serverless: Some(Box::new(ServerlessSpec { cloud: cloud_enum, - region: params.region, + region, })), pod: None, }; - Ok(CreateIndexRequest { - name: params.name, - dimension: params.dimension, + let create_index_request = CreateIndexRequest { + name, + dimension: dimension.try_into().unwrap(), metric: metric_enum, spec: Some(Box::new(create_index_request_spec)), - }) + }; + + match openapi::apis::manage_indexes_api::create_index( + &self.openapi_config(), + create_index_request, + ) + .await + { + Ok(index) => Ok(index), + Err(e) => Err(PineconeError::CreateIndexError { openapi_error: e }), + } } } @@ -74,43 +95,67 @@ mod tests { use super::*; use core::panic; use mockito::mock; - use openapi::models::serverless_spec; use serial_test::serial; use tokio; #[tokio::test] #[serial] - async fn test_create_serverless_index_req() { + async fn test_create_serverless_index() { + let _m = mock("POST", "/indexes") + .with_status(201) + .with_header("content-type", "application/json") + .with_body( + r#" + { + "name": "index_name", + "dimension": 10, + "metric": "cosine", + "host": "host1", + "spec": { + "serverless": { + "cloud": "aws", + "region": "us-east-1" + } + }, + "status": { + "ready": true, + "state": "Initializing" + } + } + "#, + ) + .create(); + let pinecone = Pinecone::new( Some("api_key".to_string()), - Some("controller_url".to_string()), + Some(mockito::server_url()), None, None, ); - let params = CreateServerlessIndexRequest { - name: "index_name".to_string(), - dimension: 10, - metric: Some("cosine".to_string()), - cloud: Some("gcp".to_string()), - region: "us-east-1".to_string(), - }; - - let create_index_request = pinecone.unwrap().create_serverless_index_req(params); + let name = "index_name".to_string(); + let dimension = 10; + + let create_index_request = pinecone.unwrap().create_serverless_index( + name, + dimension, + Metric::Cosine, + Cloud::Aws, + "us-east-1".to_string() + ).await; assert!(create_index_request.is_ok()); - let create_index_request = create_index_request.unwrap(); - assert_eq!(create_index_request.name, "index_name"); - assert_eq!(create_index_request.dimension, 10); - assert_eq!(create_index_request.metric, Some(Metric::Cosine)); + let create_index_req = create_index_request.unwrap(); + assert_eq!(create_index_req.name, "index_name"); + assert_eq!(create_index_req.dimension, 10); + assert_eq!(create_index_req.metric, openapi::models::index_model::Metric::Cosine); - let spec = create_index_request.spec.unwrap(); - let serverless_spec = spec.serverless.unwrap(); - assert_eq!(serverless_spec.cloud, serverless_spec::Cloud::Gcp); - assert_eq!(serverless_spec.region, "us-east-1"); + let spec = create_index_req.spec.serverless.unwrap(); + assert_eq!(spec.cloud, openapi::models::serverless_spec::Cloud::Aws); + assert_eq!(spec.region, "us-east-1"); } #[tokio::test] - async fn test_create_serverless_index() { + async fn test_create_index_serverless() { // Create a mock server let _m = mock("POST", "/indexes") .with_status(201) @@ -143,15 +188,17 @@ mod tests { None, None, ); - let params = CreateServerlessIndexRequest { + let params = CreateIndexParams { name: "index_name".to_string(), dimension: 10, - metric: Some("euclidean".to_string()), - cloud: Some("aws".to_string()), - region: "us-east-1".to_string(), + metric: Metric::Euclidean, + spec: Spec::Serverless { + cloud: Cloud::Aws, + region: "us-east-1".to_string(), + }, }; - let result = pinecone.unwrap().create_serverless_index(params).await; + let result = pinecone.unwrap().create_index(params).await; match result { Ok(index) => { diff --git a/pinecone_sdk/src/models/create_index_request_params.rs b/pinecone_sdk/src/models/create_index_request_params.rs index 02dede4..adba3bd 100644 --- a/pinecone_sdk/src/models/create_index_request_params.rs +++ b/pinecone_sdk/src/models/create_index_request_params.rs @@ -1,20 +1,50 @@ -pub struct CreateServerlessIndexRequest { +#[derive(Debug)] +pub struct CreateIndexParams { pub name: String, - pub dimension: i32, - pub metric: Option, - pub cloud: Option, - pub region: String, + pub dimension: u32, + pub metric: Metric, + pub spec: Spec, } -impl CreateServerlessIndexRequest { - pub fn new(name: &str, dimension: i32, metric: Option<&str>, cloud: Option<&str>, region: &str) -> CreateServerlessIndexRequest { - CreateServerlessIndexRequest { +impl CreateIndexParams { + pub fn new(name: &str, dimension: u32, metric: Option, spec: Spec) -> CreateIndexParams { + CreateIndexParams { name: name.to_string(), dimension, - metric: if let Some(metric) = metric { Some(metric.to_string()) } else { None }, - cloud: if let Some(cloud) = cloud { Some(cloud.to_string()) } else { None }, - region: region.to_string(), + metric: metric.map(|x| x).unwrap_or(Metric::Cosine), + spec, } } -} \ No newline at end of file +} + +#[derive(Debug)] +pub enum Spec { + Serverless { + cloud: Cloud, + region: String, + }, + Pod { + environment: String, + replicas: Option, + shards: Option, + pod_type: String, + pods: i32, + metadata_config: Option, + source_collection: Option, + }, +} + +#[derive(Debug)] +pub enum Metric { + Cosine, + Euclidean, + Dotproduct, +} + +#[derive(Debug)] +pub enum Cloud { + Aws, + Gcp, + Azure, +} From cd0dfcbbc277e5d45220834a891fe76fed8e3c0e Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Mon, 10 Jun 2024 10:56:18 -0400 Subject: [PATCH 09/20] add builder pattern for CreateIndexParams --- .../src/models/create_index_request_params.rs | 103 +++++++++++++++++- 1 file changed, 100 insertions(+), 3 deletions(-) diff --git a/pinecone_sdk/src/models/create_index_request_params.rs b/pinecone_sdk/src/models/create_index_request_params.rs index adba3bd..e930586 100644 --- a/pinecone_sdk/src/models/create_index_request_params.rs +++ b/pinecone_sdk/src/models/create_index_request_params.rs @@ -16,9 +16,63 @@ impl CreateIndexParams { spec, } } + + // constructs builder for CreateIndexParams + pub fn builder() -> CreateIndexParamsBuilder { + CreateIndexParamsBuilder::new() + } } -#[derive(Debug)] +pub struct CreateIndexParamsBuilder { + name: String, + dimension: u32, + metric: Option, + spec: Option, +} + +impl CreateIndexParamsBuilder { + pub fn new() -> CreateIndexParamsBuilder { + CreateIndexParamsBuilder { + name: "".to_string(), + dimension: 0, + metric: None, + spec: None, + } + } + + pub fn with_name(mut self, name: &str) -> CreateIndexParamsBuilder { + self.name = name.to_string(); + self + } + + pub fn with_dimension(mut self, dimension: u32) -> CreateIndexParamsBuilder { + // want to eventually throw an error if dimension is 0? + self.dimension = dimension; + self + } + + pub fn with_metric(mut self, metric: Metric) -> CreateIndexParamsBuilder { + self.metric = Some(metric); + self + } + + pub fn with_spec(mut self, spec: Spec) -> CreateIndexParamsBuilder { + self.spec = Some(spec); + self + } + + // constructs CreateIndexParams from CreateIndexParamsBuilder fields + pub fn build(self) -> CreateIndexParams { + CreateIndexParams { + name: self.name, + dimension: self.dimension, + metric: self.metric.map(|x| x).unwrap_or(Metric::Cosine), + spec: self.spec.map(|x| x).unwrap(), + } + } +} + +#[derive(Debug, PartialEq)] pub enum Spec { Serverless { cloud: Cloud, @@ -35,16 +89,59 @@ pub enum Spec { }, } -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub enum Metric { Cosine, Euclidean, Dotproduct, } -#[derive(Debug)] +#[derive(Debug, PartialEq)] pub enum Cloud { Aws, Gcp, Azure, } + + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_create_index_params() { + let create_index_params = CreateIndexParams::new("test_index", 10, None, Spec::Serverless { + cloud: Cloud::Aws, + region: "us-west-2".to_string(), + }); + + assert_eq!(create_index_params.name, "test_index"); + assert_eq!(create_index_params.dimension, 10); + assert_eq!(create_index_params.metric, Metric::Cosine); + assert_eq!(create_index_params.spec, Spec::Serverless { + cloud: Cloud::Aws, + region: "us-west-2".to_string(), + }); + } + + #[test] + fn test_create_index_params_builder() { + let create_index_params = CreateIndexParams::builder() + .with_name("test_index") + .with_dimension(10) + .with_metric(Metric::Cosine) + .with_spec(Spec::Serverless { + cloud: Cloud::Aws, + region: "us-west-2".to_string(), + }) + .build(); + + assert_eq!(create_index_params.name, "test_index"); + assert_eq!(create_index_params.dimension, 10); + assert_eq!(create_index_params.metric, Metric::Cosine); + assert_eq!(create_index_params.spec, Spec::Serverless { + cloud: Cloud::Aws, + region: "us-west-2".to_string(), + }); + } +} \ No newline at end of file From 901f0d58dc94b20561aee5da1967af8cd773e121 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Mon, 10 Jun 2024 14:02:09 -0400 Subject: [PATCH 10/20] added some error checking + more tests --- .../src/models/create_index_request_params.rs | 105 +++++++++++++++--- pinecone_sdk/src/utils/errors.rs | 26 +++++ 2 files changed, 114 insertions(+), 17 deletions(-) diff --git a/pinecone_sdk/src/models/create_index_request_params.rs b/pinecone_sdk/src/models/create_index_request_params.rs index e930586..2fc5633 100644 --- a/pinecone_sdk/src/models/create_index_request_params.rs +++ b/pinecone_sdk/src/models/create_index_request_params.rs @@ -1,5 +1,7 @@ +use crate::utils::errors::PineconeError; -#[derive(Debug)] + +#[derive(Debug, PartialEq)] pub struct CreateIndexParams { pub name: String, pub dimension: u32, @@ -12,7 +14,7 @@ impl CreateIndexParams { CreateIndexParams { name: name.to_string(), dimension, - metric: metric.map(|x| x).unwrap_or(Metric::Cosine), + metric: metric.unwrap_or(Metric::Cosine), spec, } } @@ -24,8 +26,8 @@ impl CreateIndexParams { } pub struct CreateIndexParamsBuilder { - name: String, - dimension: u32, + name: Option, + dimension: Option, metric: Option, spec: Option, } @@ -33,21 +35,21 @@ pub struct CreateIndexParamsBuilder { impl CreateIndexParamsBuilder { pub fn new() -> CreateIndexParamsBuilder { CreateIndexParamsBuilder { - name: "".to_string(), - dimension: 0, + name: None, + dimension: None, metric: None, spec: None, } } pub fn with_name(mut self, name: &str) -> CreateIndexParamsBuilder { - self.name = name.to_string(); + self.name = Some(name.to_string()); self } pub fn with_dimension(mut self, dimension: u32) -> CreateIndexParamsBuilder { // want to eventually throw an error if dimension is 0? - self.dimension = dimension; + self.dimension = Some(dimension); self } @@ -62,13 +64,26 @@ impl CreateIndexParamsBuilder { } // constructs CreateIndexParams from CreateIndexParamsBuilder fields - pub fn build(self) -> CreateIndexParams { - CreateIndexParams { - name: self.name, - dimension: self.dimension, - metric: self.metric.map(|x| x).unwrap_or(Metric::Cosine), - spec: self.spec.map(|x| x).unwrap(), - } + pub fn build(self) -> Result { + let name = match self.name { + Some(name) => name, + None => Err(PineconeError::MissingNameError)?, + }; + let dimension = match self.dimension { + Some(dimension) => dimension, + None => Err(PineconeError::MissingDimensionError)?, + }; + let spec = match self.spec { + Some(spec) => spec, + None => Err(PineconeError::MissingSpecError)?, + }; + + Ok(CreateIndexParams { + name, + dimension, + metric: self.metric.unwrap_or(Metric::Cosine), + spec, + }) } } @@ -134,7 +149,29 @@ mod tests { cloud: Cloud::Aws, region: "us-west-2".to_string(), }) - .build(); + .build() + .unwrap(); + + assert_eq!(create_index_params.name, "test_index"); + assert_eq!(create_index_params.dimension, 10); + assert_eq!(create_index_params.metric, Metric::Cosine); + assert_eq!(create_index_params.spec, Spec::Serverless { + cloud: Cloud::Aws, + region: "us-west-2".to_string(), + }); + } + + #[test] + fn test_builder_missing_metric() { + let create_index_params = CreateIndexParams::builder() + .with_name("test_index") + .with_dimension(10) + .with_spec(Spec::Serverless { + cloud: Cloud::Aws, + region: "us-west-2".to_string(), + }) + .build() + .unwrap(); assert_eq!(create_index_params.name, "test_index"); assert_eq!(create_index_params.dimension, 10); @@ -144,4 +181,38 @@ mod tests { region: "us-west-2".to_string(), }); } -} \ No newline at end of file + + #[test] + fn test_missing_name() { + let create_index_params = CreateIndexParams::builder() + .with_dimension(10) + .with_metric(Metric::Cosine) + .build(); + + assert!(create_index_params.is_err()); + assert_eq!(create_index_params, Err(PineconeError::MissingNameError)); + } + + #[test] + fn test_missing_dimension() { + let create_index_params = CreateIndexParams::builder() + .with_name("test_index") + .with_metric(Metric::Cosine) + .build(); + + assert!(create_index_params.is_err()); + assert_eq!(create_index_params, Err(PineconeError::MissingDimensionError)); + } + + #[test] + fn test_missing_spec() { + let create_index_params = CreateIndexParams::builder() + .with_name("test_index") + .with_dimension(10) + .build(); + + assert!(create_index_params.is_err()); + assert_eq!(create_index_params, Err(PineconeError::MissingSpecError)); + } + +} diff --git a/pinecone_sdk/src/utils/errors.rs b/pinecone_sdk/src/utils/errors.rs index 639b3dd..c426277 100644 --- a/pinecone_sdk/src/utils/errors.rs +++ b/pinecone_sdk/src/utils/errors.rs @@ -11,6 +11,15 @@ pub enum PineconeError { openapi_error: OpenAPIError, }, + #[snafu(display("Index name missing."))] + MissingNameError, + + #[snafu(display("Dimension missing."))] + MissingDimensionError, + + #[snafu(display("Spec missing."))] + MissingSpecError, + #[snafu(display("Failed to parse headers: {}", json_error))] InvalidHeadersError { json_error: serde_json::Error }, @@ -23,3 +32,20 @@ pub enum PineconeError { #[snafu(display("Invalid region."))] InvalidRegionError, } + +impl PartialEq for PineconeError { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (PineconeError::APIKeyMissingError, PineconeError::APIKeyMissingError) => true, + (PineconeError::CreateIndexError { .. }, PineconeError::CreateIndexError { .. }) => true, + (PineconeError::MissingNameError, PineconeError::MissingNameError) => true, + (PineconeError::MissingDimensionError, PineconeError::MissingDimensionError) => true, + (PineconeError::MissingSpecError, PineconeError::MissingSpecError) => true, + (PineconeError::InvalidHeadersError { .. }, PineconeError::InvalidHeadersError { .. }) => true, + (PineconeError::InvalidCloudError { .. }, PineconeError::InvalidCloudError { .. }) => true, + (PineconeError::InvalidMetricError { .. }, PineconeError::InvalidMetricError { .. }) => true, + (PineconeError::InvalidRegionError, PineconeError::InvalidRegionError) => true, + _ => false, + } + } +} \ No newline at end of file From 1bc27f12aa0d38ae0658dd7f1819723f67babde0 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Mon, 10 Jun 2024 14:19:51 -0400 Subject: [PATCH 11/20] added builder and more tests for Pinecone --- .../src/models/create_index_request_params.rs | 25 +++--- pinecone_sdk/src/pinecone.rs | 82 ++++++++++++++++++- 2 files changed, 93 insertions(+), 14 deletions(-) diff --git a/pinecone_sdk/src/models/create_index_request_params.rs b/pinecone_sdk/src/models/create_index_request_params.rs index 2fc5633..badc6a4 100644 --- a/pinecone_sdk/src/models/create_index_request_params.rs +++ b/pinecone_sdk/src/models/create_index_request_params.rs @@ -122,9 +122,10 @@ pub enum Cloud { #[cfg(test)] mod tests { use super::*; + use tokio; - #[test] - fn test_create_index_params() { + #[tokio::test] + async fn test_create_index_params() { let create_index_params = CreateIndexParams::new("test_index", 10, None, Spec::Serverless { cloud: Cloud::Aws, region: "us-west-2".to_string(), @@ -139,8 +140,8 @@ mod tests { }); } - #[test] - fn test_create_index_params_builder() { + #[tokio::test] + async fn test_create_index_params_builder() { let create_index_params = CreateIndexParams::builder() .with_name("test_index") .with_dimension(10) @@ -161,8 +162,8 @@ mod tests { }); } - #[test] - fn test_builder_missing_metric() { + #[tokio::test] + async fn test_builder_missing_metric() { let create_index_params = CreateIndexParams::builder() .with_name("test_index") .with_dimension(10) @@ -182,8 +183,8 @@ mod tests { }); } - #[test] - fn test_missing_name() { + #[tokio::test] + async fn test_missing_name() { let create_index_params = CreateIndexParams::builder() .with_dimension(10) .with_metric(Metric::Cosine) @@ -193,8 +194,8 @@ mod tests { assert_eq!(create_index_params, Err(PineconeError::MissingNameError)); } - #[test] - fn test_missing_dimension() { + #[tokio::test] + async fn test_missing_dimension() { let create_index_params = CreateIndexParams::builder() .with_name("test_index") .with_metric(Metric::Cosine) @@ -204,8 +205,8 @@ mod tests { assert_eq!(create_index_params, Err(PineconeError::MissingDimensionError)); } - #[test] - fn test_missing_spec() { + #[tokio::test] + async fn test_missing_spec() { let create_index_params = CreateIndexParams::builder() .with_name("test_index") .with_dimension(10) diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index e52c044..754ac89 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -73,11 +73,65 @@ impl Pinecone { }) } + pub fn builder() -> PineconeBuilder { + PineconeBuilder::new() + } + pub fn openapi_config(&self) -> &Configuration { &self.openapi_config } } +pub struct PineconeBuilder { + api_key: Option, + control_plane_host: Option, + additional_headers: Option>, + source_tag: Option, +} + +impl PineconeBuilder { + pub fn new() -> PineconeBuilder { + PineconeBuilder { + api_key: None, + control_plane_host: None, + additional_headers: None, + source_tag: None, + } + } + + pub fn with_api_key(mut self, api_key: &str) -> PineconeBuilder { + self.api_key = Some(api_key.to_string()); + self + } + + pub fn with_control_plane_host(mut self, control_plane_host: &str) -> PineconeBuilder { + self.control_plane_host = Some(control_plane_host.to_string()); + self + } + + pub fn with_additional_headers( + mut self, + additional_headers: HashMap, + ) -> PineconeBuilder { + self.additional_headers = Some(additional_headers); + self + } + + pub fn with_source_tag(mut self, source_tag: &str) -> PineconeBuilder { + self.source_tag = Some(source_tag.to_string()); + self + } + + pub fn build(self) -> Result { + Pinecone::new( + self.api_key, + self.control_plane_host, + self.additional_headers, + self.source_tag, + ) + } +} + #[cfg(test)] mod tests { use std::env; @@ -147,10 +201,10 @@ mod tests { ); assert!(pinecone.is_err()); - assert!(matches!( + assert_eq!( pinecone.err().unwrap(), PineconeError::APIKeyMissingError - )); + ); } #[tokio::test] @@ -344,6 +398,30 @@ mod tests { pinecone.as_ref().unwrap().config.additional_headers, mock_arg_headers.clone() ); + } + + #[tokio::test] + async fn test_builder() { + let pinecone = Pinecone::builder() + .with_api_key("mock-api-key") + .build(); + assert_eq!(pinecone.unwrap().config.api_key, "mock-api-key"); + } + + #[tokio::test] + async fn test_builder_all_params() { + let pinecone = Pinecone::builder() + .with_api_key("mock-api-key") + .with_additional_headers(HashMap::from([("header1".to_string(), "value1".to_string())])) + .with_control_plane_host("mock-controller-host") + .with_source_tag("mock-source-tag") + .build() + .unwrap(); + + assert_eq!(pinecone.config.api_key, "mock-api-key"); + assert_eq!(pinecone.config.additional_headers, HashMap::from([("header1".to_string(), "value1".to_string())])); + assert_eq!(pinecone.config.controller_url, "mock-controller-host"); + assert_eq!(pinecone.config.source_tag, Some("mock-source-tag".to_string())); } } From 91c885c7f4967e85f5575b1a185cd70be6bf040b Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Mon, 10 Jun 2024 14:34:40 -0400 Subject: [PATCH 12/20] added some comments --- pinecone_sdk/src/pinecone.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index 754ac89..df28357 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -73,6 +73,7 @@ impl Pinecone { }) } + // constructs a PineconeBuilder instance pub fn builder() -> PineconeBuilder { PineconeBuilder::new() } @@ -122,6 +123,7 @@ impl PineconeBuilder { self } + // Constructs Pinecone instance from PineconeBuilder fields pub fn build(self) -> Result { Pinecone::new( self.api_key, From 099b0d938b9ae4faa1c3e612d47c9890de2cf79b Mon Sep 17 00:00:00 2001 From: Emily Yu Date: Mon, 10 Jun 2024 15:09:37 -0400 Subject: [PATCH 13/20] Update params builder required arg validation --- .../src/models/create_index_request_params.rs | 83 +++++++++++-------- 1 file changed, 47 insertions(+), 36 deletions(-) diff --git a/pinecone_sdk/src/models/create_index_request_params.rs b/pinecone_sdk/src/models/create_index_request_params.rs index badc6a4..a82f10b 100644 --- a/pinecone_sdk/src/models/create_index_request_params.rs +++ b/pinecone_sdk/src/models/create_index_request_params.rs @@ -1,6 +1,5 @@ use crate::utils::errors::PineconeError; - #[derive(Debug, PartialEq)] pub struct CreateIndexParams { pub name: String, @@ -10,7 +9,12 @@ pub struct CreateIndexParams { } impl CreateIndexParams { - pub fn new(name: &str, dimension: u32, metric: Option, spec: Spec) -> CreateIndexParams { + pub fn new( + name: &str, + dimension: u32, + metric: Option, + spec: Spec, + ) -> CreateIndexParams { CreateIndexParams { name: name.to_string(), dimension, @@ -65,18 +69,10 @@ impl CreateIndexParamsBuilder { // constructs CreateIndexParams from CreateIndexParamsBuilder fields pub fn build(self) -> Result { - let name = match self.name { - Some(name) => name, - None => Err(PineconeError::MissingNameError)?, - }; - let dimension = match self.dimension { - Some(dimension) => dimension, - None => Err(PineconeError::MissingDimensionError)?, - }; - let spec = match self.spec { - Some(spec) => spec, - None => Err(PineconeError::MissingSpecError)?, - }; + // required parameters + let name = self.name.ok_or(PineconeError::MissingNameError)?; + let dimension = self.dimension.ok_or(PineconeError::MissingDimensionError)?; + let spec = self.spec.ok_or(PineconeError::MissingSpecError)?; Ok(CreateIndexParams { name, @@ -118,7 +114,6 @@ pub enum Cloud { Azure, } - #[cfg(test)] mod tests { use super::*; @@ -126,18 +121,26 @@ mod tests { #[tokio::test] async fn test_create_index_params() { - let create_index_params = CreateIndexParams::new("test_index", 10, None, Spec::Serverless { - cloud: Cloud::Aws, - region: "us-west-2".to_string(), - }); + let create_index_params = CreateIndexParams::new( + "test_index", + 10, + None, + Spec::Serverless { + cloud: Cloud::Aws, + region: "us-west-2".to_string(), + }, + ); assert_eq!(create_index_params.name, "test_index"); assert_eq!(create_index_params.dimension, 10); assert_eq!(create_index_params.metric, Metric::Cosine); - assert_eq!(create_index_params.spec, Spec::Serverless { - cloud: Cloud::Aws, - region: "us-west-2".to_string(), - }); + assert_eq!( + create_index_params.spec, + Spec::Serverless { + cloud: Cloud::Aws, + region: "us-west-2".to_string(), + } + ); } #[tokio::test] @@ -156,10 +159,13 @@ mod tests { assert_eq!(create_index_params.name, "test_index"); assert_eq!(create_index_params.dimension, 10); assert_eq!(create_index_params.metric, Metric::Cosine); - assert_eq!(create_index_params.spec, Spec::Serverless { - cloud: Cloud::Aws, - region: "us-west-2".to_string(), - }); + assert_eq!( + create_index_params.spec, + Spec::Serverless { + cloud: Cloud::Aws, + region: "us-west-2".to_string(), + } + ); } #[tokio::test] @@ -177,10 +183,13 @@ mod tests { assert_eq!(create_index_params.name, "test_index"); assert_eq!(create_index_params.dimension, 10); assert_eq!(create_index_params.metric, Metric::Cosine); - assert_eq!(create_index_params.spec, Spec::Serverless { - cloud: Cloud::Aws, - region: "us-west-2".to_string(), - }); + assert_eq!( + create_index_params.spec, + Spec::Serverless { + cloud: Cloud::Aws, + region: "us-west-2".to_string(), + } + ); } #[tokio::test] @@ -189,7 +198,7 @@ mod tests { .with_dimension(10) .with_metric(Metric::Cosine) .build(); - + assert!(create_index_params.is_err()); assert_eq!(create_index_params, Err(PineconeError::MissingNameError)); } @@ -200,9 +209,12 @@ mod tests { .with_name("test_index") .with_metric(Metric::Cosine) .build(); - + assert!(create_index_params.is_err()); - assert_eq!(create_index_params, Err(PineconeError::MissingDimensionError)); + assert_eq!( + create_index_params, + Err(PineconeError::MissingDimensionError) + ); } #[tokio::test] @@ -211,9 +223,8 @@ mod tests { .with_name("test_index") .with_dimension(10) .build(); - + assert!(create_index_params.is_err()); assert_eq!(create_index_params, Err(PineconeError::MissingSpecError)); } - } From 86a9db2143cec95b9cdb726fd6c6549a6f046130 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Fri, 7 Jun 2024 10:28:19 -0400 Subject: [PATCH 14/20] added integration test --- .github/workflows/ci.yaml | 4 +++- pinecone_sdk/tests/integration_test.rs | 23 +++++++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) create mode 100644 pinecone_sdk/tests/integration_test.rs diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 3fd67d6..bbcdd4b 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -20,4 +20,6 @@ jobs: - name: Build run: cargo build - name: Run tests - run: cargo test --verbose \ No newline at end of file + env: + PINECONE_API_KEY: ${{ secrets.PINECONE_API_KEY }} + run: cargo test --verbose diff --git a/pinecone_sdk/tests/integration_test.rs b/pinecone_sdk/tests/integration_test.rs new file mode 100644 index 0000000..ed6a94e --- /dev/null +++ b/pinecone_sdk/tests/integration_test.rs @@ -0,0 +1,23 @@ +use pinecone_sdk::pinecone::Pinecone; + +#[tokio::test] +async fn test_list_serverless_index() { + let pinecone = Pinecone::new("b41b6453-9756-45aa-8d8d-a51c295d3c78".to_string(), None); + let list_response = pinecone.list_indexes().await; + + match list_response { + Ok(index_list) => { + match index_list.indexes { + Some(indexes) => { + assert_eq!(indexes.len(), 0); + } + None => { + assert!(false, "Expected indexes to be Some"); + } + } + } + Err(e) => { + panic!("Error: {:?}", e); + } + } +} From bde2ba9b7624c243df3a0c23da43fbfb7a739235 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Fri, 7 Jun 2024 10:41:42 -0400 Subject: [PATCH 15/20] fixed syntax in ci.yaml --- .github/workflows/ci.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index bbcdd4b..d42d865 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -20,6 +20,6 @@ jobs: - name: Build run: cargo build - name: Run tests - env: - PINECONE_API_KEY: ${{ secrets.PINECONE_API_KEY }} + env: + PINECONE_API_KEY: ${{ secrets.PINECONE_API_KEY }} run: cargo test --verbose From df10f7bc09be50d282bb8f09faa02efd5de8ebfd Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Fri, 7 Jun 2024 11:47:35 -0400 Subject: [PATCH 16/20] added one integration test --- pinecone_sdk/tests/integration_tests.rs | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 pinecone_sdk/tests/integration_tests.rs diff --git a/pinecone_sdk/tests/integration_tests.rs b/pinecone_sdk/tests/integration_tests.rs new file mode 100644 index 0000000..fc764ba --- /dev/null +++ b/pinecone_sdk/tests/integration_tests.rs @@ -0,0 +1,22 @@ +use pinecone_sdk::pinecone::Pinecone; +use std::env; + +#[tokio::test] +async fn test_list_serverless_index() { + let api_key = env::var("PINECONE_API_KEY").unwrap(); + let pinecone = Pinecone::new(api_key.to_string(), None); + let list_response = pinecone.list_indexes().await; + + match list_response { + Ok(index_list) => { + if let Some(indexes) = index_list.indexes { + assert_eq!(indexes.len(), 0); + } else { + panic!("No indexes found"); + } + }, + Err(e) => { + panic!("Error: {:?}", e); + } + } +} From bba2160e236bdb95d4051add19cdb41c1812f582 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Mon, 10 Jun 2024 15:29:29 -0400 Subject: [PATCH 17/20] rebased to create-index --- pinecone_sdk/src/utils/errors.rs | 2 +- pinecone_sdk/tests/integration_test.rs | 5 ++++- pinecone_sdk/tests/integration_tests.rs | 22 ---------------------- 3 files changed, 5 insertions(+), 24 deletions(-) delete mode 100644 pinecone_sdk/tests/integration_tests.rs diff --git a/pinecone_sdk/src/utils/errors.rs b/pinecone_sdk/src/utils/errors.rs index c426277..9049e83 100644 --- a/pinecone_sdk/src/utils/errors.rs +++ b/pinecone_sdk/src/utils/errors.rs @@ -48,4 +48,4 @@ impl PartialEq for PineconeError { _ => false, } } -} \ No newline at end of file +} diff --git a/pinecone_sdk/tests/integration_test.rs b/pinecone_sdk/tests/integration_test.rs index ed6a94e..f920de3 100644 --- a/pinecone_sdk/tests/integration_test.rs +++ b/pinecone_sdk/tests/integration_test.rs @@ -2,7 +2,10 @@ use pinecone_sdk::pinecone::Pinecone; #[tokio::test] async fn test_list_serverless_index() { - let pinecone = Pinecone::new("b41b6453-9756-45aa-8d8d-a51c295d3c78".to_string(), None); + let pinecone = Pinecone::builder() + .with_api_key("b41b6453-9756-45aa-8d8d-a51c295d3c78") + .build() + .unwrap(); let list_response = pinecone.list_indexes().await; match list_response { diff --git a/pinecone_sdk/tests/integration_tests.rs b/pinecone_sdk/tests/integration_tests.rs deleted file mode 100644 index fc764ba..0000000 --- a/pinecone_sdk/tests/integration_tests.rs +++ /dev/null @@ -1,22 +0,0 @@ -use pinecone_sdk::pinecone::Pinecone; -use std::env; - -#[tokio::test] -async fn test_list_serverless_index() { - let api_key = env::var("PINECONE_API_KEY").unwrap(); - let pinecone = Pinecone::new(api_key.to_string(), None); - let list_response = pinecone.list_indexes().await; - - match list_response { - Ok(index_list) => { - if let Some(indexes) = index_list.indexes { - assert_eq!(indexes.len(), 0); - } else { - panic!("No indexes found"); - } - }, - Err(e) => { - panic!("Error: {:?}", e); - } - } -} From 6998156a49443bddb4e0a48c74c4572d2c3e5c11 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Mon, 10 Jun 2024 16:32:43 -0400 Subject: [PATCH 18/20] rebased to main --- pinecone_sdk/src/config.rs | 2 +- pinecone_sdk/src/control/create_index.rs | 222 ----------------- pinecone_sdk/src/control/list_indexes.rs | 5 +- pinecone_sdk/src/control/mod.rs | 5 +- pinecone_sdk/src/lib.rs | 3 +- .../src/models/create_index_request_params.rs | 230 ------------------ pinecone_sdk/src/models/mod.rs | 1 - pinecone_sdk/src/pinecone.rs | 90 +------ pinecone_sdk/src/utils/errors.rs | 43 +--- pinecone_sdk/src/utils/mod.rs | 2 +- pinecone_sdk/tests/integration_test.rs | 27 +- 11 files changed, 38 insertions(+), 592 deletions(-) delete mode 100644 pinecone_sdk/src/control/create_index.rs delete mode 100644 pinecone_sdk/src/models/create_index_request_params.rs delete mode 100644 pinecone_sdk/src/models/mod.rs diff --git a/pinecone_sdk/src/config.rs b/pinecone_sdk/src/config.rs index 9bfea0e..18b11a3 100644 --- a/pinecone_sdk/src/config.rs +++ b/pinecone_sdk/src/config.rs @@ -17,4 +17,4 @@ impl Config { source_tag, } } -} +} \ No newline at end of file diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs deleted file mode 100644 index e0ab7a6..0000000 --- a/pinecone_sdk/src/control/create_index.rs +++ /dev/null @@ -1,222 +0,0 @@ -use crate::models::create_index_request_params::{CreateIndexParams, Spec, Metric, Cloud}; -use crate::pinecone::Pinecone; -use crate::utils::errors::PineconeError; -use openapi::models::create_index_request; -use openapi::models::serverless_spec; -use openapi::models::{CreateIndexRequest, CreateIndexRequestSpec, IndexModel, ServerlessSpec}; - -impl Pinecone { - // Creates a new index based on the provided parameters - pub async fn create_index( - &self, - params: CreateIndexParams - ) -> Result { - - // Check if creating serverless or pod-based index and call respective builder function - match params.spec { - Spec::Serverless { cloud, region } => { - self.create_serverless_index( - params.name, - params.dimension, - params.metric, - cloud, - region - ).await - } - Spec::Pod { - environment: _, - replicas: _, - shards: _, - pod_type: _, - pods: _, - metadata_config: _, - source_collection: _, - } => { - // eventually change this to be pod index - self.create_serverless_index(params.name, params.dimension, params.metric, Cloud::Aws, "".to_string()).await - } - } - } - - // Creates serverless index - async fn create_serverless_index( - &self, - name: String, - dimension: u32, - metric: Metric, - cloud: Cloud, - region: String - ) -> Result { - - // clean metric enum - let metric_enum = match metric { - Metric::Cosine => Some(create_index_request::Metric::Cosine), - Metric::Dotproduct => Some(create_index_request::Metric::Dotproduct), - Metric::Euclidean => Some(create_index_request::Metric::Euclidean), - }; - - // clean cloud enum - let cloud_enum = match cloud { - Cloud::Aws => serverless_spec::Cloud::Aws, - Cloud::Gcp => serverless_spec::Cloud::Gcp, - Cloud::Azure => serverless_spec::Cloud::Azure, - }; - - // create request specs - let create_index_request_spec = CreateIndexRequestSpec { - serverless: Some(Box::new(ServerlessSpec { - cloud: cloud_enum, - region, - })), - pod: None, - }; - - let create_index_request = CreateIndexRequest { - name, - dimension: dimension.try_into().unwrap(), - metric: metric_enum, - spec: Some(Box::new(create_index_request_spec)), - }; - - match openapi::apis::manage_indexes_api::create_index( - &self.openapi_config(), - create_index_request, - ) - .await - { - Ok(index) => Ok(index), - Err(e) => Err(PineconeError::CreateIndexError { openapi_error: e }), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use core::panic; - use mockito::mock; - use serial_test::serial; - use tokio; - - #[tokio::test] - #[serial] - async fn test_create_serverless_index() { - let _m = mock("POST", "/indexes") - .with_status(201) - .with_header("content-type", "application/json") - .with_body( - r#" - { - "name": "index_name", - "dimension": 10, - "metric": "cosine", - "host": "host1", - "spec": { - "serverless": { - "cloud": "aws", - "region": "us-east-1" - } - }, - "status": { - "ready": true, - "state": "Initializing" - } - } - "#, - ) - .create(); - - let pinecone = Pinecone::new( - Some("api_key".to_string()), - Some(mockito::server_url()), - None, - None, - ); - let name = "index_name".to_string(); - let dimension = 10; - - let create_index_request = pinecone.unwrap().create_serverless_index( - name, - dimension, - Metric::Cosine, - Cloud::Aws, - "us-east-1".to_string() - ).await; - assert!(create_index_request.is_ok()); - - let create_index_req = create_index_request.unwrap(); - assert_eq!(create_index_req.name, "index_name"); - assert_eq!(create_index_req.dimension, 10); - assert_eq!(create_index_req.metric, openapi::models::index_model::Metric::Cosine); - - let spec = create_index_req.spec.serverless.unwrap(); - assert_eq!(spec.cloud, openapi::models::serverless_spec::Cloud::Aws); - assert_eq!(spec.region, "us-east-1"); - } - - #[tokio::test] - async fn test_create_index_serverless() { - // Create a mock server - let _m = mock("POST", "/indexes") - .with_status(201) - .with_header("content-type", "application/json") - .with_body( - r#" - { - "name": "index_name", - "dimension": 10, - "metric": "euclidean", - "host": "host1", - "spec": { - "serverless": { - "cloud": "aws", - "region": "us-east-1" - } - }, - "status": { - "ready": true, - "state": "Initializing" - } - } - "#, - ) - .create(); - - let pinecone = Pinecone::new( - Some("api_key".to_string()), - Some(mockito::server_url()), - None, - None, - ); - let params = CreateIndexParams { - name: "index_name".to_string(), - dimension: 10, - metric: Metric::Euclidean, - spec: Spec::Serverless { - cloud: Cloud::Aws, - region: "us-east-1".to_string(), - }, - }; - - let result = pinecone.unwrap().create_index(params).await; - - match result { - Ok(index) => { - assert_eq!(index.name, "index_name"); - assert_eq!(index.dimension, 10); - assert_eq!( - index.metric, - openapi::models::index_model::Metric::Euclidean - ); - let spec = *index.spec; - let serverless_spec = spec.serverless.unwrap(); - assert_eq!( - serverless_spec.cloud, - openapi::models::serverless_spec::Cloud::Aws - ); - assert_eq!(serverless_spec.region, "us-east-1"); - } - Err(e) => panic!("{}", e), - } - } -} diff --git a/pinecone_sdk/src/control/list_indexes.rs b/pinecone_sdk/src/control/list_indexes.rs index 4cea7ba..8ff707f 100644 --- a/pinecone_sdk/src/control/list_indexes.rs +++ b/pinecone_sdk/src/control/list_indexes.rs @@ -15,8 +15,11 @@ impl Pinecone { #[cfg(test)] mod tests { use super::*; + use crate::config::Config; use crate::control::list_indexes::models::index_model::Metric; use mockito::mock; + use openapi::apis::configuration::ApiKey; + use openapi::apis::configuration::Configuration; use openapi::models::IndexList; use openapi::models::IndexModel; use tokio; @@ -95,4 +98,4 @@ mod tests { Err(err) => panic!("Expected Ok, got Err: {:?}", err), } } -} +} \ No newline at end of file diff --git a/pinecone_sdk/src/control/mod.rs b/pinecone_sdk/src/control/mod.rs index 5b4c3c3..61b19f0 100644 --- a/pinecone_sdk/src/control/mod.rs +++ b/pinecone_sdk/src/control/mod.rs @@ -1,2 +1,3 @@ -pub mod create_index; -pub mod list_indexes; +mod list_indexes; + +pub use list_indexes::*; \ No newline at end of file diff --git a/pinecone_sdk/src/lib.rs b/pinecone_sdk/src/lib.rs index df1d5a8..84c7bc3 100644 --- a/pinecone_sdk/src/lib.rs +++ b/pinecone_sdk/src/lib.rs @@ -1,5 +1,4 @@ pub mod config; pub mod control; pub mod pinecone; -pub mod utils; -pub mod models; +pub mod utils; \ No newline at end of file diff --git a/pinecone_sdk/src/models/create_index_request_params.rs b/pinecone_sdk/src/models/create_index_request_params.rs deleted file mode 100644 index a82f10b..0000000 --- a/pinecone_sdk/src/models/create_index_request_params.rs +++ /dev/null @@ -1,230 +0,0 @@ -use crate::utils::errors::PineconeError; - -#[derive(Debug, PartialEq)] -pub struct CreateIndexParams { - pub name: String, - pub dimension: u32, - pub metric: Metric, - pub spec: Spec, -} - -impl CreateIndexParams { - pub fn new( - name: &str, - dimension: u32, - metric: Option, - spec: Spec, - ) -> CreateIndexParams { - CreateIndexParams { - name: name.to_string(), - dimension, - metric: metric.unwrap_or(Metric::Cosine), - spec, - } - } - - // constructs builder for CreateIndexParams - pub fn builder() -> CreateIndexParamsBuilder { - CreateIndexParamsBuilder::new() - } -} - -pub struct CreateIndexParamsBuilder { - name: Option, - dimension: Option, - metric: Option, - spec: Option, -} - -impl CreateIndexParamsBuilder { - pub fn new() -> CreateIndexParamsBuilder { - CreateIndexParamsBuilder { - name: None, - dimension: None, - metric: None, - spec: None, - } - } - - pub fn with_name(mut self, name: &str) -> CreateIndexParamsBuilder { - self.name = Some(name.to_string()); - self - } - - pub fn with_dimension(mut self, dimension: u32) -> CreateIndexParamsBuilder { - // want to eventually throw an error if dimension is 0? - self.dimension = Some(dimension); - self - } - - pub fn with_metric(mut self, metric: Metric) -> CreateIndexParamsBuilder { - self.metric = Some(metric); - self - } - - pub fn with_spec(mut self, spec: Spec) -> CreateIndexParamsBuilder { - self.spec = Some(spec); - self - } - - // constructs CreateIndexParams from CreateIndexParamsBuilder fields - pub fn build(self) -> Result { - // required parameters - let name = self.name.ok_or(PineconeError::MissingNameError)?; - let dimension = self.dimension.ok_or(PineconeError::MissingDimensionError)?; - let spec = self.spec.ok_or(PineconeError::MissingSpecError)?; - - Ok(CreateIndexParams { - name, - dimension, - metric: self.metric.unwrap_or(Metric::Cosine), - spec, - }) - } -} - -#[derive(Debug, PartialEq)] -pub enum Spec { - Serverless { - cloud: Cloud, - region: String, - }, - Pod { - environment: String, - replicas: Option, - shards: Option, - pod_type: String, - pods: i32, - metadata_config: Option, - source_collection: Option, - }, -} - -#[derive(Debug, PartialEq)] -pub enum Metric { - Cosine, - Euclidean, - Dotproduct, -} - -#[derive(Debug, PartialEq)] -pub enum Cloud { - Aws, - Gcp, - Azure, -} - -#[cfg(test)] -mod tests { - use super::*; - use tokio; - - #[tokio::test] - async fn test_create_index_params() { - let create_index_params = CreateIndexParams::new( - "test_index", - 10, - None, - Spec::Serverless { - cloud: Cloud::Aws, - region: "us-west-2".to_string(), - }, - ); - - assert_eq!(create_index_params.name, "test_index"); - assert_eq!(create_index_params.dimension, 10); - assert_eq!(create_index_params.metric, Metric::Cosine); - assert_eq!( - create_index_params.spec, - Spec::Serverless { - cloud: Cloud::Aws, - region: "us-west-2".to_string(), - } - ); - } - - #[tokio::test] - async fn test_create_index_params_builder() { - let create_index_params = CreateIndexParams::builder() - .with_name("test_index") - .with_dimension(10) - .with_metric(Metric::Cosine) - .with_spec(Spec::Serverless { - cloud: Cloud::Aws, - region: "us-west-2".to_string(), - }) - .build() - .unwrap(); - - assert_eq!(create_index_params.name, "test_index"); - assert_eq!(create_index_params.dimension, 10); - assert_eq!(create_index_params.metric, Metric::Cosine); - assert_eq!( - create_index_params.spec, - Spec::Serverless { - cloud: Cloud::Aws, - region: "us-west-2".to_string(), - } - ); - } - - #[tokio::test] - async fn test_builder_missing_metric() { - let create_index_params = CreateIndexParams::builder() - .with_name("test_index") - .with_dimension(10) - .with_spec(Spec::Serverless { - cloud: Cloud::Aws, - region: "us-west-2".to_string(), - }) - .build() - .unwrap(); - - assert_eq!(create_index_params.name, "test_index"); - assert_eq!(create_index_params.dimension, 10); - assert_eq!(create_index_params.metric, Metric::Cosine); - assert_eq!( - create_index_params.spec, - Spec::Serverless { - cloud: Cloud::Aws, - region: "us-west-2".to_string(), - } - ); - } - - #[tokio::test] - async fn test_missing_name() { - let create_index_params = CreateIndexParams::builder() - .with_dimension(10) - .with_metric(Metric::Cosine) - .build(); - - assert!(create_index_params.is_err()); - assert_eq!(create_index_params, Err(PineconeError::MissingNameError)); - } - - #[tokio::test] - async fn test_missing_dimension() { - let create_index_params = CreateIndexParams::builder() - .with_name("test_index") - .with_metric(Metric::Cosine) - .build(); - - assert!(create_index_params.is_err()); - assert_eq!( - create_index_params, - Err(PineconeError::MissingDimensionError) - ); - } - - #[tokio::test] - async fn test_missing_spec() { - let create_index_params = CreateIndexParams::builder() - .with_name("test_index") - .with_dimension(10) - .build(); - - assert!(create_index_params.is_err()); - assert_eq!(create_index_params, Err(PineconeError::MissingSpecError)); - } -} diff --git a/pinecone_sdk/src/models/mod.rs b/pinecone_sdk/src/models/mod.rs deleted file mode 100644 index b1fe893..0000000 --- a/pinecone_sdk/src/models/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod create_index_request_params; \ No newline at end of file diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index df28357..0305931 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -54,7 +54,7 @@ impl Pinecone { additional_headers, source_tag, }; - + let user_agent = get_user_agent(&config); let openapi_config = Configuration { @@ -73,67 +73,11 @@ impl Pinecone { }) } - // constructs a PineconeBuilder instance - pub fn builder() -> PineconeBuilder { - PineconeBuilder::new() - } - pub fn openapi_config(&self) -> &Configuration { &self.openapi_config } } -pub struct PineconeBuilder { - api_key: Option, - control_plane_host: Option, - additional_headers: Option>, - source_tag: Option, -} - -impl PineconeBuilder { - pub fn new() -> PineconeBuilder { - PineconeBuilder { - api_key: None, - control_plane_host: None, - additional_headers: None, - source_tag: None, - } - } - - pub fn with_api_key(mut self, api_key: &str) -> PineconeBuilder { - self.api_key = Some(api_key.to_string()); - self - } - - pub fn with_control_plane_host(mut self, control_plane_host: &str) -> PineconeBuilder { - self.control_plane_host = Some(control_plane_host.to_string()); - self - } - - pub fn with_additional_headers( - mut self, - additional_headers: HashMap, - ) -> PineconeBuilder { - self.additional_headers = Some(additional_headers); - self - } - - pub fn with_source_tag(mut self, source_tag: &str) -> PineconeBuilder { - self.source_tag = Some(source_tag.to_string()); - self - } - - // Constructs Pinecone instance from PineconeBuilder fields - pub fn build(self) -> Result { - Pinecone::new( - self.api_key, - self.control_plane_host, - self.additional_headers, - self.source_tag, - ) - } -} - #[cfg(test)] mod tests { use std::env; @@ -203,10 +147,10 @@ mod tests { ); assert!(pinecone.is_err()); - assert_eq!( + assert!(matches!( pinecone.err().unwrap(), PineconeError::APIKeyMissingError - ); + )); } #[tokio::test] @@ -332,7 +276,6 @@ mod tests { pinecone.err().unwrap(), PineconeError::InvalidHeadersError { .. } )); - remove_env_var("PINECONE_ADDITIONAL_HEADERS"); } #[tokio::test] @@ -401,29 +344,4 @@ mod tests { mock_arg_headers.clone() ); } - - #[tokio::test] - async fn test_builder() { - let pinecone = Pinecone::builder() - .with_api_key("mock-api-key") - .build(); - - assert_eq!(pinecone.unwrap().config.api_key, "mock-api-key"); - } - - #[tokio::test] - async fn test_builder_all_params() { - let pinecone = Pinecone::builder() - .with_api_key("mock-api-key") - .with_additional_headers(HashMap::from([("header1".to_string(), "value1".to_string())])) - .with_control_plane_host("mock-controller-host") - .with_source_tag("mock-source-tag") - .build() - .unwrap(); - - assert_eq!(pinecone.config.api_key, "mock-api-key"); - assert_eq!(pinecone.config.additional_headers, HashMap::from([("header1".to_string(), "value1".to_string())])); - assert_eq!(pinecone.config.controller_url, "mock-controller-host"); - assert_eq!(pinecone.config.source_tag, Some("mock-source-tag".to_string())); - } -} +} \ No newline at end of file diff --git a/pinecone_sdk/src/utils/errors.rs b/pinecone_sdk/src/utils/errors.rs index 9049e83..69024d6 100644 --- a/pinecone_sdk/src/utils/errors.rs +++ b/pinecone_sdk/src/utils/errors.rs @@ -1,4 +1,3 @@ -use openapi::apis::{manage_indexes_api::CreateIndexError, Error as OpenAPIError}; use snafu::prelude::*; #[derive(Debug, Snafu)] @@ -6,46 +5,6 @@ pub enum PineconeError { #[snafu(display("API key missing."))] APIKeyMissingError, - #[snafu(display("API key missing."))] - CreateIndexError { - openapi_error: OpenAPIError, - }, - - #[snafu(display("Index name missing."))] - MissingNameError, - - #[snafu(display("Dimension missing."))] - MissingDimensionError, - - #[snafu(display("Spec missing."))] - MissingSpecError, - #[snafu(display("Failed to parse headers: {}", json_error))] InvalidHeadersError { json_error: serde_json::Error }, - - #[snafu(display("Invalid cloud '{}'.", cloud))] - InvalidCloudError { cloud: String }, - - #[snafu(display("Invalid metric '{}'.", metric))] - InvalidMetricError { metric: String }, - - #[snafu(display("Invalid region."))] - InvalidRegionError, -} - -impl PartialEq for PineconeError { - fn eq(&self, other: &Self) -> bool { - match (self, other) { - (PineconeError::APIKeyMissingError, PineconeError::APIKeyMissingError) => true, - (PineconeError::CreateIndexError { .. }, PineconeError::CreateIndexError { .. }) => true, - (PineconeError::MissingNameError, PineconeError::MissingNameError) => true, - (PineconeError::MissingDimensionError, PineconeError::MissingDimensionError) => true, - (PineconeError::MissingSpecError, PineconeError::MissingSpecError) => true, - (PineconeError::InvalidHeadersError { .. }, PineconeError::InvalidHeadersError { .. }) => true, - (PineconeError::InvalidCloudError { .. }, PineconeError::InvalidCloudError { .. }) => true, - (PineconeError::InvalidMetricError { .. }, PineconeError::InvalidMetricError { .. }) => true, - (PineconeError::InvalidRegionError, PineconeError::InvalidRegionError) => true, - _ => false, - } - } -} +} \ No newline at end of file diff --git a/pinecone_sdk/src/utils/mod.rs b/pinecone_sdk/src/utils/mod.rs index 1eaf948..8549526 100644 --- a/pinecone_sdk/src/utils/mod.rs +++ b/pinecone_sdk/src/utils/mod.rs @@ -1,2 +1,2 @@ pub mod errors; -pub mod user_agent; +pub mod user_agent; \ No newline at end of file diff --git a/pinecone_sdk/tests/integration_test.rs b/pinecone_sdk/tests/integration_test.rs index f920de3..98e697e 100644 --- a/pinecone_sdk/tests/integration_test.rs +++ b/pinecone_sdk/tests/integration_test.rs @@ -2,10 +2,7 @@ use pinecone_sdk::pinecone::Pinecone; #[tokio::test] async fn test_list_serverless_index() { - let pinecone = Pinecone::builder() - .with_api_key("b41b6453-9756-45aa-8d8d-a51c295d3c78") - .build() - .unwrap(); + let pinecone = Pinecone::new(Some("b41b6453-9756-45aa-8d8d-a51c295d3c78".to_string()), None, None, None).unwrap(); let list_response = pinecone.list_indexes().await; match list_response { @@ -24,3 +21,25 @@ async fn test_list_serverless_index() { } } } + +#[tokio::test] +async fn test_list_serverless_index_env() { + let pinecone = Pinecone::new(None, None, None, None).unwrap(); + let list_response = pinecone.list_indexes().await; + + match list_response { + Ok(index_list) => { + match index_list.indexes { + Some(indexes) => { + assert_eq!(indexes.len(), 0); + } + None => { + assert!(false, "Expected indexes to be Some"); + } + } + } + Err(e) => { + panic!("Error: {:?}", e); + } + } +} \ No newline at end of file From 832391928c9a251eb94bc16f5f8303e991043490 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Mon, 10 Jun 2024 16:46:15 -0400 Subject: [PATCH 19/20] edited integration test asserts --- pinecone_sdk/tests/integration_test.rs | 32 ++------------------------ 1 file changed, 2 insertions(+), 30 deletions(-) diff --git a/pinecone_sdk/tests/integration_test.rs b/pinecone_sdk/tests/integration_test.rs index 98e697e..09ef230 100644 --- a/pinecone_sdk/tests/integration_test.rs +++ b/pinecone_sdk/tests/integration_test.rs @@ -5,21 +5,7 @@ async fn test_list_serverless_index() { let pinecone = Pinecone::new(Some("b41b6453-9756-45aa-8d8d-a51c295d3c78".to_string()), None, None, None).unwrap(); let list_response = pinecone.list_indexes().await; - match list_response { - Ok(index_list) => { - match index_list.indexes { - Some(indexes) => { - assert_eq!(indexes.len(), 0); - } - None => { - assert!(false, "Expected indexes to be Some"); - } - } - } - Err(e) => { - panic!("Error: {:?}", e); - } - } + assert!(list_response.is_ok()); } #[tokio::test] @@ -27,19 +13,5 @@ async fn test_list_serverless_index_env() { let pinecone = Pinecone::new(None, None, None, None).unwrap(); let list_response = pinecone.list_indexes().await; - match list_response { - Ok(index_list) => { - match index_list.indexes { - Some(indexes) => { - assert_eq!(indexes.len(), 0); - } - None => { - assert!(false, "Expected indexes to be Some"); - } - } - } - Err(e) => { - panic!("Error: {:?}", e); - } - } + assert!(list_response.is_ok()); } \ No newline at end of file From 0de0350a38e7b65096d974f7c0cd525341870ee5 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Mon, 10 Jun 2024 17:56:46 -0400 Subject: [PATCH 20/20] deleted test --- pinecone_sdk/tests/integration_test.rs | 8 -------- 1 file changed, 8 deletions(-) diff --git a/pinecone_sdk/tests/integration_test.rs b/pinecone_sdk/tests/integration_test.rs index 09ef230..ae970a9 100644 --- a/pinecone_sdk/tests/integration_test.rs +++ b/pinecone_sdk/tests/integration_test.rs @@ -1,13 +1,5 @@ use pinecone_sdk::pinecone::Pinecone; -#[tokio::test] -async fn test_list_serverless_index() { - let pinecone = Pinecone::new(Some("b41b6453-9756-45aa-8d8d-a51c295d3c78".to_string()), None, None, None).unwrap(); - let list_response = pinecone.list_indexes().await; - - assert!(list_response.is_ok()); -} - #[tokio::test] async fn test_list_serverless_index_env() { let pinecone = Pinecone::new(None, None, None, None).unwrap();