From e6d95374a2553d3aff01d032fad40a9443908ad9 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Tue, 4 Jun 2024 16:15:14 -0400 Subject: [PATCH 01/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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/28] 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 a84154ca7d52d0480750dc91e952ad42e49efd01 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Tue, 11 Jun 2024 13:28:58 -0400 Subject: [PATCH 14/28] remove builder for CreateIndexParams --- .../src/models/create_index_request_params.rs | 145 ------------------ 1 file changed, 145 deletions(-) diff --git a/pinecone_sdk/src/models/create_index_request_params.rs b/pinecone_sdk/src/models/create_index_request_params.rs index a82f10b..a05a487 100644 --- a/pinecone_sdk/src/models/create_index_request_params.rs +++ b/pinecone_sdk/src/models/create_index_request_params.rs @@ -1,4 +1,3 @@ -use crate::utils::errors::PineconeError; #[derive(Debug, PartialEq)] pub struct CreateIndexParams { @@ -22,65 +21,6 @@ impl CreateIndexParams { 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)] @@ -142,89 +82,4 @@ mod tests { } ); } - - #[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)); - } } From 8c289ea14c4503862be085436d7cb91906519574 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Tue, 11 Jun 2024 14:47:01 -0400 Subject: [PATCH 15/28] change builder function names --- pinecone_sdk/src/pinecone.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index d36ea87..44eb89e 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -100,17 +100,17 @@ impl PineconeBuilder { } } - pub fn with_api_key(mut self, api_key: &str) -> PineconeBuilder { + pub fn 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 { + pub fn 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( + pub fn additional_headers( mut self, additional_headers: HashMap, ) -> PineconeBuilder { @@ -118,7 +118,7 @@ impl PineconeBuilder { self } - pub fn with_source_tag(mut self, source_tag: &str) -> PineconeBuilder { + pub fn source_tag(mut self, source_tag: &str) -> PineconeBuilder { self.source_tag = Some(source_tag.to_string()); self } From 3c62d01dc80d974e33506c327b6cc9db44fe918a Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Tue, 11 Jun 2024 14:49:16 -0400 Subject: [PATCH 16/28] fixed test cases --- pinecone_sdk/src/pinecone.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index 44eb89e..1e619ab 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -405,7 +405,7 @@ mod tests { #[tokio::test] async fn test_builder() { let pinecone = Pinecone::builder() - .with_api_key("mock-api-key") + .api_key("mock-api-key") .build(); assert_eq!(pinecone.unwrap().config.api_key, "mock-api-key"); @@ -414,10 +414,10 @@ mod tests { #[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") + .api_key("mock-api-key") + .additional_headers(HashMap::from([("header1".to_string(), "value1".to_string())])) + .control_plane_host("mock-controller-host") + .source_tag("mock-source-tag") .build() .unwrap(); From 7b647889bb85b4ae117aacb1aa570d5db40146fa Mon Sep 17 00:00:00 2001 From: Emily Yu Date: Wed, 12 Jun 2024 10:30:07 -0400 Subject: [PATCH 17/28] Implement Into trait to convert to openapi types --- pinecone_sdk/src/control/create_index.rs | 21 +++++------------- .../src/models/create_index_request_params.rs | 22 +++++++++++++++++++ 2 files changed, 27 insertions(+), 16 deletions(-) diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs index 2f99774..870ee39 100644 --- a/pinecone_sdk/src/control/create_index.rs +++ b/pinecone_sdk/src/control/create_index.rs @@ -47,25 +47,14 @@ impl Pinecone { 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, - }; + // convert to openapi types + let openapi_metric: create_index_request::Metric = metric.into(); + let openapi_cloud: serverless_spec::Cloud = cloud.into(); // create request specs let create_index_request_spec = CreateIndexRequestSpec { serverless: Some(Box::new(ServerlessSpec { - cloud: cloud_enum, + cloud: openapi_cloud, region, })), pod: None, @@ -74,7 +63,7 @@ impl Pinecone { let create_index_request = CreateIndexRequest { name, dimension: dimension.try_into().unwrap(), - metric: metric_enum, + metric: Some(openapi_metric), spec: Some(Box::new(create_index_request_spec)), }; diff --git a/pinecone_sdk/src/models/create_index_request_params.rs b/pinecone_sdk/src/models/create_index_request_params.rs index a05a487..c62187a 100644 --- a/pinecone_sdk/src/models/create_index_request_params.rs +++ b/pinecone_sdk/src/models/create_index_request_params.rs @@ -1,3 +1,5 @@ +use openapi::models::create_index_request; +use openapi::models::serverless_spec; #[derive(Debug, PartialEq)] pub struct CreateIndexParams { @@ -47,6 +49,16 @@ pub enum Metric { Dotproduct, } +impl Into for Metric { + fn into(self) -> create_index_request::Metric { + match self { + Metric::Cosine => create_index_request::Metric::Cosine, + Metric::Euclidean => create_index_request::Metric::Euclidean, + Metric::Dotproduct => create_index_request::Metric::Dotproduct, + } + } +} + #[derive(Debug, PartialEq)] pub enum Cloud { Aws, @@ -54,6 +66,16 @@ pub enum Cloud { Azure, } +impl Into for Cloud { + fn into(self) -> serverless_spec::Cloud { + match self { + Cloud::Aws => serverless_spec::Cloud::Aws, + Cloud::Gcp => serverless_spec::Cloud::Gcp, + Cloud::Azure => serverless_spec::Cloud::Azure, + } + } +} + #[cfg(test)] mod tests { use super::*; From ef7870f661dc462e9bf4e6acc35a282bb9f6f7aa Mon Sep 17 00:00:00 2001 From: Emily Yu Date: Wed, 12 Jun 2024 10:51:04 -0400 Subject: [PATCH 18/28] Add missing documentation --- pinecone_sdk/src/control/create_index.rs | 2 +- pinecone_sdk/src/control/mod.rs | 3 ++ pinecone_sdk/src/lib.rs | 2 +- .../src/models/create_index_request_params.rs | 32 +++++++++++++++++++ pinecone_sdk/src/models/mod.rs | 3 +- pinecone_sdk/src/pinecone.rs | 10 ++++-- pinecone_sdk/src/utils/errors.rs | 21 ++++++++---- 7 files changed, 62 insertions(+), 11 deletions(-) diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs index 870ee39..7f7d67c 100644 --- a/pinecone_sdk/src/control/create_index.rs +++ b/pinecone_sdk/src/control/create_index.rs @@ -6,7 +6,7 @@ use openapi::models::serverless_spec; use openapi::models::{CreateIndexRequest, CreateIndexRequestSpec, IndexModel, ServerlessSpec}; impl Pinecone { - // Creates a new index based on the provided parameters + /// Creates a new index based on the provided parameters pub async fn create_index( &self, params: CreateIndexParams diff --git a/pinecone_sdk/src/control/mod.rs b/pinecone_sdk/src/control/mod.rs index 5b4c3c3..7c509b4 100644 --- a/pinecone_sdk/src/control/mod.rs +++ b/pinecone_sdk/src/control/mod.rs @@ -1,2 +1,5 @@ +/// Module for creating an index. pub mod create_index; + +/// Module for listing all indexes. pub mod list_indexes; diff --git a/pinecone_sdk/src/lib.rs b/pinecone_sdk/src/lib.rs index d67ee84..79499e4 100644 --- a/pinecone_sdk/src/lib.rs +++ b/pinecone_sdk/src/lib.rs @@ -18,5 +18,5 @@ pub mod pinecone; /// Utility modules. pub mod utils; -// Structure models. +/// Structure models. pub mod models; diff --git a/pinecone_sdk/src/models/create_index_request_params.rs b/pinecone_sdk/src/models/create_index_request_params.rs index c62187a..59670dc 100644 --- a/pinecone_sdk/src/models/create_index_request_params.rs +++ b/pinecone_sdk/src/models/create_index_request_params.rs @@ -1,15 +1,23 @@ use openapi::models::create_index_request; use openapi::models::serverless_spec; +/// Request parameters for creating an index. #[derive(Debug, PartialEq)] pub struct CreateIndexParams { + /// The name of the index. + /// Resource name must be 1-45 characters long, start and end with an alphanumeric character, and consist only of + /// lower case alphanumeric characters or '-'. pub name: String, + /// The dimensions of the vectors to be inserted in the index. pub dimension: u32, + /// The distance metric to be used for similarity search. You can use 'euclidean', 'cosine', or 'dotproduct'. pub metric: Metric, + /// The spec object defines how the index should be deployed. pub spec: Spec, } impl CreateIndexParams { + /// Create a new instance of CreateIndexParams. pub fn new( name: &str, dimension: u32, @@ -25,27 +33,47 @@ impl CreateIndexParams { } } +/// Spec for the index. +/// Can be either serverless or pod. #[derive(Debug, PartialEq)] pub enum Spec { + /// Spec for a serverless index. Serverless { + /// The public cloud where you would like your index hosted. cloud: Cloud, + /// The region where you would like your index to be created. region: String, }, + /// Spec for a pod-based index. Pod { + /// The environment where the index is hosted. environment: String, + /// The number of replicas. Replicas duplicate your index. They provide higher availability and throughput. + /// Replicas can be scaled up or down as your needs change. replicas: Option, + /// The number of shards. Shards split your data across multiple pods so you can fit more data into an index. shards: Option, + /// The type of pod to use. One of `s1`, `p1`, or `p2` appended with `.` and one of `x1`, `x2`, `x4`, or `x8`. pod_type: String, + /// The number of pods to be used in the index. This should be equal to `shards` x `replicas`. pods: i32, + /// Configuration for the behavior of Pinecone's internal metadata index. By default, all metadata is indexed; + /// when `metadata_config` is present, only specified metadata fields are indexed. + /// These configurations are only valid for use with pod-based indexes. metadata_config: Option, + /// The name of the collection to be used as the source for the index. source_collection: Option, }, } +/// Distance metric to be used for similarity search. #[derive(Debug, PartialEq)] pub enum Metric { + /// Cosine similarity Cosine, + /// Euclidean distance Euclidean, + /// Dot product Dotproduct, } @@ -59,10 +87,14 @@ impl Into for Metric { } } +/// Cloud where the index should be hosted. #[derive(Debug, PartialEq)] pub enum Cloud { + /// Amazon Web Services Aws, + /// Google Cloud Platform Gcp, + /// Microsoft Azure Azure, } diff --git a/pinecone_sdk/src/models/mod.rs b/pinecone_sdk/src/models/mod.rs index b1fe893..53a7c78 100644 --- a/pinecone_sdk/src/models/mod.rs +++ b/pinecone_sdk/src/models/mod.rs @@ -1 +1,2 @@ -pub mod create_index_request_params; \ No newline at end of file +/// Request parameters for create_index. +pub mod create_index_request_params; diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index 2b33fb7..60b6b01 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -95,7 +95,7 @@ impl Pinecone { }) } - // constructs a PineconeBuilder instance + /// Constructs a PineconeBuilder instance pub fn builder() -> PineconeBuilder { PineconeBuilder::new() } @@ -106,6 +106,7 @@ impl Pinecone { } } +/// The `PineconeBuilder` struct is used to construct a `Pinecone` instance with a builder pattern. pub struct PineconeBuilder { api_key: Option, control_plane_host: Option, @@ -114,6 +115,7 @@ pub struct PineconeBuilder { } impl PineconeBuilder { + /// Constructs a new PineconeBuilder instance. pub fn new() -> PineconeBuilder { PineconeBuilder { api_key: None, @@ -123,16 +125,19 @@ impl PineconeBuilder { } } + /// Sets the API key for the Pinecone instance. pub fn api_key(mut self, api_key: &str) -> PineconeBuilder { self.api_key = Some(api_key.to_string()); self } + /// Sets the control plane host for the Pinecone instance. pub fn control_plane_host(mut self, control_plane_host: &str) -> PineconeBuilder { self.control_plane_host = Some(control_plane_host.to_string()); self } + /// Sets additional headers for the Pinecone instance. pub fn additional_headers( mut self, additional_headers: HashMap, @@ -141,12 +146,13 @@ impl PineconeBuilder { self } + /// Sets the source tag for the Pinecone instance. pub fn source_tag(mut self, source_tag: &str) -> PineconeBuilder { self.source_tag = Some(source_tag.to_string()); self } - // Constructs Pinecone instance from PineconeBuilder fields + /// Constructs Pinecone instance from PineconeBuilder fields. pub fn build(self) -> Result { Pinecone::new( self.api_key, diff --git a/pinecone_sdk/src/utils/errors.rs b/pinecone_sdk/src/utils/errors.rs index c20f8ad..08ce7a0 100644 --- a/pinecone_sdk/src/utils/errors.rs +++ b/pinecone_sdk/src/utils/errors.rs @@ -8,13 +8,19 @@ pub enum PineconeError { #[snafu(display("API key missing."))] APIKeyMissingError, + /// CreateIndexError: Failed to create an index. #[snafu(display("API key missing."))] CreateIndexError { + /// openapi_error: Error object for OpenAPI error. openapi_error: OpenAPIError, }, + /// InvalidCloudError: Provided cloud is not valid. #[snafu(display("Invalid cloud '{}'.", cloud))] - InvalidCloudError { cloud: String }, + InvalidCloudError { + /// cloud: Cloud name. + cloud: String, + }, /// InvalidHeadersError: Provided headers are not valid. Expects JSON. #[snafu(display("Failed to parse headers: {}", json_error))] @@ -23,18 +29,22 @@ pub enum PineconeError { json_error: serde_json::Error, }, + /// InvalidMetricError: Provided metric is not valid. #[snafu(display("Invalid metric '{}'.", metric))] - InvalidMetricError { metric: String }, - - #[snafu(display("Invalid region."))] - InvalidRegionError, + InvalidMetricError { + /// metric: Metric name. + metric: String, + }, + /// MissingNameError: Index name is missing. #[snafu(display("Index name missing."))] MissingNameError, + /// MissingDimensionError: Index dimension is missing. #[snafu(display("Dimension missing."))] MissingDimensionError, + /// MissingSpecError: Index spec is missing. #[snafu(display("Spec missing."))] MissingSpecError, } @@ -60,7 +70,6 @@ impl PartialEq for PineconeError { PineconeError::InvalidMetricError { .. }, PineconeError::InvalidMetricError { .. }, ) => true, - (PineconeError::InvalidRegionError, PineconeError::InvalidRegionError) => true, _ => false, } } From a40712f8d1966fa7a40a9e1b9ff158691d230c88 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Thu, 13 Jun 2024 09:39:45 -0400 Subject: [PATCH 19/28] changed Pinecone to PineconeClient --- pinecone_sdk/src/control/create_index.rs | 8 ++--- pinecone_sdk/src/control/list_indexes.rs | 10 +++--- pinecone_sdk/src/pinecone.rs | 40 ++++++++++++------------ pinecone_sdk/tests/integration_test.rs | 4 +-- 4 files changed, 31 insertions(+), 31 deletions(-) diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs index 7f7d67c..5f4e1c8 100644 --- a/pinecone_sdk/src/control/create_index.rs +++ b/pinecone_sdk/src/control/create_index.rs @@ -1,11 +1,11 @@ use crate::models::create_index_request_params::{CreateIndexParams, Spec, Metric, Cloud}; -use crate::pinecone::Pinecone; +use crate::pinecone::PineconeClient; 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 { +impl PineconeClient { /// Creates a new index based on the provided parameters pub async fn create_index( &self, @@ -113,7 +113,7 @@ mod tests { ) .create(); - let pinecone = Pinecone::new( + let pinecone = PineconeClient::new( Some("api_key".to_string()), Some(mockito::server_url()), None, @@ -169,7 +169,7 @@ mod tests { ) .create(); - let pinecone = Pinecone::new( + let pinecone = PineconeClient::new( Some("api_key".to_string()), Some(mockito::server_url()), None, diff --git a/pinecone_sdk/src/control/list_indexes.rs b/pinecone_sdk/src/control/list_indexes.rs index d24f76e..990ffc7 100644 --- a/pinecone_sdk/src/control/list_indexes.rs +++ b/pinecone_sdk/src/control/list_indexes.rs @@ -1,10 +1,10 @@ -use crate::pinecone::Pinecone; +use crate::pinecone::PineconeClient; use openapi::apis::manage_indexes_api; use openapi::apis::manage_indexes_api::ListIndexesError; use openapi::apis::Error; use openapi::models; -impl Pinecone { +impl PineconeClient { /// Lists all indexes. /// /// The results include a description of all indexes in your project, including the @@ -17,12 +17,12 @@ impl Pinecone { /// ### Example /// /// ``` - /// # use pinecone_sdk::pinecone::Pinecone; + /// # use pinecone_sdk::pinecone::PineconeClient; /// # use pinecone_sdk::utils::errors::PineconeError; /// # #[tokio::main] /// # async fn main() -> Result<(), PineconeError>{ /// # // Create a Pinecone client with the API key and controller host. - /// # let pinecone = Pinecone::new(None, None, None, None).unwrap(); + /// # let pinecone = PineconeClient::new(None, None, None, None).unwrap(); /// # /// // List all indexes in the project. /// let index_list = pinecone.list_indexes(); @@ -87,7 +87,7 @@ mod tests { // Construct Pinecone instance with the mock server URL let api_key = "test_api_key".to_string(); - let pinecone = Pinecone::new(Some(api_key), Some(mockito::server_url()), None, None) + let pinecone = PineconeClient::new(Some(api_key), Some(mockito::server_url()), None, None) .expect("Failed to create Pinecone instance"); // Call list_indexes and verify the result diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index 60b6b01..b800e10 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -8,7 +8,7 @@ use std::collections::HashMap; /// The `Pinecone` struct is the main entry point for interacting with Pinecone via this Rust SDK. #[derive(Debug, Clone)] -pub struct Pinecone { +pub struct PineconeClient { /// Configuration for the Pinecone SDK struct. config: Config, @@ -16,7 +16,7 @@ pub struct Pinecone { openapi_config: Configuration, } -impl Pinecone { +impl PineconeClient { /// The `Pinecone` struct is the main entry point for interacting with Pinecone via this Rust SDK. /// It is used to create, delete, and manage your indexes and collections. /// @@ -30,10 +30,10 @@ impl Pinecone { /// ### Example /// /// ``` - /// use pinecone_sdk::pinecone::Pinecone; + /// use pinecone_sdk::pinecone::PineconeClient; /// /// // Create a Pinecone client with the API key and controller host. - /// let pinecone = Pinecone::new(Some("INSERT_API_KEY".to_string()), Some("INSERT_CONTROLLER_HOST".to_string()), None, None); + /// let pinecone = PineconeClient::new(Some("INSERT_API_KEY".to_string()), Some("INSERT_CONTROLLER_HOST".to_string()), None, None); /// ``` pub fn new( api_key: Option, @@ -89,7 +89,7 @@ impl Pinecone { ..Default::default() }; - Ok(Pinecone { + Ok(PineconeClient { config, openapi_config, }) @@ -153,8 +153,8 @@ impl PineconeBuilder { } /// Constructs Pinecone instance from PineconeBuilder fields. - pub fn build(self) -> Result { - Pinecone::new( + pub fn build(self) -> Result { + PineconeClient::new( self.api_key, self.control_plane_host, self.additional_headers, @@ -173,7 +173,7 @@ mod tests { let mock_api_key = "mock-arg-api-key".to_string(); let mock_controller_host = "mock-arg-controller-host".to_string(); - let pinecone = Pinecone::new( + let pinecone = PineconeClient::new( Some(mock_api_key.clone()), Some(mock_controller_host.clone()), Some(HashMap::new()), @@ -192,7 +192,7 @@ mod tests { let mock_controller_host = "mock-arg-controller-host".to_string(); temp_env::with_var("PINECONE_API_KEY", Some(mock_api_key.as_str()), || { - let pinecone = Pinecone::new( + let pinecone = PineconeClient::new( None, Some(mock_controller_host.clone()), Some(HashMap::new()), @@ -211,7 +211,7 @@ mod tests { let mock_controller_host = "mock-arg-controller-host".to_string(); temp_env::with_var_unset("PINECONE_API_KEY", || { - let pinecone = Pinecone::new( + let pinecone = PineconeClient::new( None, Some(mock_controller_host.clone()), Some(HashMap::new()), @@ -229,7 +229,7 @@ mod tests { async fn test_arg_host() -> Result<(), PineconeError> { let mock_api_key = "mock-arg-api-key".to_string(); let mock_controller_host = "mock-arg-controller-host".to_string(); - let pinecone = Pinecone::new( + let pinecone = PineconeClient::new( Some(mock_api_key.clone()), Some(mock_controller_host.clone()), Some(HashMap::new()), @@ -252,7 +252,7 @@ mod tests { Some(mock_controller_host.as_str()), || { let pinecone = - Pinecone::new(Some(mock_api_key.clone()), None, Some(HashMap::new()), None) + PineconeClient::new(Some(mock_api_key.clone()), None, Some(HashMap::new()), None) .expect("Expected to successfully create Pinecone instance with env host"); assert_eq!(pinecone.config.controller_url, mock_controller_host.clone()); @@ -267,7 +267,7 @@ mod tests { let mock_api_key = "mock-arg-api-key".to_string(); temp_env::with_var_unset("PINECONE_CONTROLLER_HOST", || { - let pinecone = Pinecone::new( + let pinecone = PineconeClient::new( Some(mock_api_key.clone()), None, Some(HashMap::new()), @@ -295,7 +295,7 @@ mod tests { ("argheader2".to_string(), "value2".to_string()), ]); - let pinecone = Pinecone::new( + let pinecone = PineconeClient::new( Some(mock_api_key.clone()), Some(mock_controller_host.clone()), Some(mock_headers.clone()), @@ -321,7 +321,7 @@ mod tests { "PINECONE_ADDITIONAL_HEADERS", Some(serde_json::to_string(&mock_headers).unwrap().as_str()), || { - let pinecone = Pinecone::new( + let pinecone = PineconeClient::new( Some(mock_api_key.clone()), Some(mock_controller_host.clone()), None, @@ -342,7 +342,7 @@ mod tests { let mock_controller_host = "mock-arg-controller-host".to_string(); temp_env::with_var("PINECONE_ADDITIONAL_HEADERS", Some("invalid-json"), || { - let pinecone = Pinecone::new( + let pinecone = PineconeClient::new( Some(mock_api_key.clone()), Some(mock_controller_host.clone()), None, @@ -365,7 +365,7 @@ mod tests { let mock_controller_host = "mock-arg-controller-host".to_string(); temp_env::with_var_unset("PINECONE_ADDITIONAL_HEADERS", || { - let pinecone = Pinecone::new( + let pinecone = PineconeClient::new( Some(mock_api_key.clone()), Some(mock_controller_host.clone()), Some(HashMap::new()), @@ -407,7 +407,7 @@ mod tests { ), ], || { - let pinecone = Pinecone::new( + let pinecone = PineconeClient::new( Some(mock_arg_api_key.clone()), Some(mock_arg_controller_host.clone()), Some(mock_arg_headers.clone()), @@ -429,14 +429,14 @@ mod tests { #[tokio::test] async fn test_builder() { - let pinecone = Pinecone::builder().api_key("mock-api-key").build(); + let pinecone = PineconeClient::builder().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() + let pinecone = PineconeClient::builder() .api_key("mock-api-key") .additional_headers(HashMap::from([( "header1".to_string(), diff --git a/pinecone_sdk/tests/integration_test.rs b/pinecone_sdk/tests/integration_test.rs index e539b81..6fe9b1c 100644 --- a/pinecone_sdk/tests/integration_test.rs +++ b/pinecone_sdk/tests/integration_test.rs @@ -1,9 +1,9 @@ -use pinecone_sdk::pinecone::Pinecone; +use pinecone_sdk::pinecone::PineconeClient; use pinecone_sdk::utils::errors::PineconeError; #[tokio::test] async fn test_list_index_env() -> Result<(), PineconeError> { - let pinecone = Pinecone::new(None, None, None, None).unwrap(); + let pinecone = PineconeClient::new(None, None, None, None).unwrap(); let _ = pinecone .list_indexes() .await From 357ea5ad9355e4fb950c3c6cda62b38e3d128e7c Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Thu, 13 Jun 2024 10:04:45 -0400 Subject: [PATCH 20/28] flattened PineconeClient fields, removed builder --- pinecone_sdk/src/control/list_indexes.rs | 2 +- pinecone_sdk/src/pinecone.rs | 177 ++++++----------------- 2 files changed, 44 insertions(+), 135 deletions(-) diff --git a/pinecone_sdk/src/control/list_indexes.rs b/pinecone_sdk/src/control/list_indexes.rs index 990ffc7..4b465d8 100644 --- a/pinecone_sdk/src/control/list_indexes.rs +++ b/pinecone_sdk/src/control/list_indexes.rs @@ -31,7 +31,7 @@ impl PineconeClient { /// ``` pub async fn list_indexes(&self) -> Result> { - let response = manage_indexes_api::list_indexes(self.openapi_config()).await?; + let response = manage_indexes_api::list_indexes(&self.openapi_config()).await?; println!("{:?}", response); Ok(response) } diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index b800e10..76a1283 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -6,18 +6,18 @@ use openapi::apis::configuration::Configuration; use serde_json; use std::collections::HashMap; -/// The `Pinecone` struct is the main entry point for interacting with Pinecone via this Rust SDK. +/// The `PineconeClient` struct is the main entry point for interacting with Pinecone via this Rust SDK. #[derive(Debug, Clone)] pub struct PineconeClient { - /// Configuration for the Pinecone SDK struct. - config: Config, - - /// OpenAPI configuration object. - openapi_config: Configuration, + api_key: String, + controller_url: String, + additional_headers: HashMap, + source_tag: Option, + user_agent: Option, } impl PineconeClient { - /// The `Pinecone` struct is the main entry point for interacting with Pinecone via this Rust SDK. + /// The `PineconeClient` struct is the main entry point for interacting with Pinecone via this Rust SDK. /// It is used to create, delete, and manage your indexes and collections. /// /// ### Configuration with environment variables @@ -42,14 +42,14 @@ impl PineconeClient { source_tag: Option, ) -> Result { // get api key - let api_key = match api_key { + let api_key_option = match api_key { Some(key) => key, None => match std::env::var("PINECONE_API_KEY") { Ok(key) => key, Err(_) => { return Err(PineconeError::APIKeyMissingError); } - }, + } }; let controller_host = control_plane_host.unwrap_or( @@ -70,97 +70,30 @@ impl PineconeClient { }, }; - let config = Config { - api_key: api_key.clone(), - controller_url: controller_host.clone(), + let config = Config::new(api_key_option.clone(), source_tag.clone()); + let user_agent = get_user_agent(&config); + + Ok(PineconeClient { + api_key: api_key_option, + controller_url: controller_host, additional_headers, source_tag, - }; - - let user_agent = get_user_agent(&config); - - let openapi_config = Configuration { - base_path: controller_host, user_agent: Some(user_agent), - api_key: Some(ApiKey { - prefix: None, - key: api_key, - }), - ..Default::default() - }; - - Ok(PineconeClient { - config, - openapi_config, }) } - /// Constructs a PineconeBuilder instance - pub fn builder() -> PineconeBuilder { - PineconeBuilder::new() - } - /// Returns the OpenAPI configuration object. - pub fn openapi_config(&self) -> &Configuration { - &self.openapi_config - } -} - -/// The `PineconeBuilder` struct is used to construct a `Pinecone` instance with a builder pattern. -pub struct PineconeBuilder { - api_key: Option, - control_plane_host: Option, - additional_headers: Option>, - source_tag: Option, -} - -impl PineconeBuilder { - /// Constructs a new PineconeBuilder instance. - pub fn new() -> PineconeBuilder { - PineconeBuilder { - api_key: None, - control_plane_host: None, - additional_headers: None, - source_tag: None, + pub fn openapi_config(&self) -> Configuration { + Configuration { + base_path: self.controller_url.clone(), + user_agent: self.user_agent.clone(), + api_key: Some(ApiKey { + prefix: None, + key: self.api_key.clone(), + }), + ..Default::default() } } - - /// Sets the API key for the Pinecone instance. - pub fn api_key(mut self, api_key: &str) -> PineconeBuilder { - self.api_key = Some(api_key.to_string()); - self - } - - /// Sets the control plane host for the Pinecone instance. - pub fn control_plane_host(mut self, control_plane_host: &str) -> PineconeBuilder { - self.control_plane_host = Some(control_plane_host.to_string()); - self - } - - /// Sets additional headers for the Pinecone instance. - pub fn additional_headers( - mut self, - additional_headers: HashMap, - ) -> PineconeBuilder { - self.additional_headers = Some(additional_headers); - self - } - - /// Sets the source tag for the Pinecone instance. - pub fn 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 { - PineconeClient::new( - self.api_key, - self.control_plane_host, - self.additional_headers, - self.source_tag, - ) - } } #[cfg(test)] @@ -181,7 +114,11 @@ mod tests { ) .expect("Expected to successfully create Pinecone instance"); - assert_eq!(pinecone.config.api_key, mock_api_key.clone()); + assert_eq!(pinecone.api_key, mock_api_key); + assert_eq!(pinecone.controller_url, mock_controller_host); + assert_eq!(pinecone.additional_headers, HashMap::new()); + assert_eq!(pinecone.source_tag, None); + assert_eq!(pinecone.user_agent, Some("lang=rust; pinecone-rust-client=0.1.0".to_string())); Ok(()) } @@ -200,7 +137,11 @@ mod tests { ) .expect("Expected to successfully create Pinecone instance"); - assert_eq!(pinecone.config.api_key, mock_api_key.clone()); + assert_eq!(pinecone.api_key, mock_api_key); + assert_eq!(pinecone.controller_url, mock_controller_host); + assert_eq!(pinecone.additional_headers, HashMap::new()); + assert_eq!(pinecone.source_tag, None); + assert_eq!(pinecone.user_agent, Some("lang=rust; pinecone-rust-client=0.1.0".to_string())); }); Ok(()) @@ -237,7 +178,7 @@ mod tests { ) .expect("Expected to successfully create Pinecone instance"); - assert_eq!(pinecone.config.controller_url, mock_controller_host.clone()); + assert_eq!(pinecone.controller_url, mock_controller_host); Ok(()) } @@ -255,7 +196,7 @@ mod tests { PineconeClient::new(Some(mock_api_key.clone()), None, Some(HashMap::new()), None) .expect("Expected to successfully create Pinecone instance with env host"); - assert_eq!(pinecone.config.controller_url, mock_controller_host.clone()); + assert_eq!(pinecone.controller_url, mock_controller_host); }, ); @@ -278,7 +219,7 @@ mod tests { ); assert_eq!( - pinecone.config.controller_url, + pinecone.controller_url, "https://api.pinecone.io".to_string() ); }); @@ -303,7 +244,7 @@ mod tests { ) .expect("Expected to successfully create Pinecone instance"); - assert_eq!(pinecone.config.additional_headers, mock_headers.clone()); + assert_eq!(pinecone.additional_headers, mock_headers); Ok(()) } @@ -329,7 +270,7 @@ mod tests { ) .expect("Expected to successfully create Pinecone instance with env headers"); - assert_eq!(pinecone.config.additional_headers, mock_headers.clone()); + assert_eq!(pinecone.additional_headers, mock_headers); }, ); @@ -373,7 +314,7 @@ mod tests { ) .expect("Expected to successfully create Pinecone instance"); - assert_eq!(pinecone.config.additional_headers, HashMap::new()); + assert_eq!(pinecone.additional_headers, HashMap::new()); }); Ok(()) @@ -415,47 +356,15 @@ mod tests { ) .expect("Expected to successfully create Pinecone instance"); - assert_eq!(pinecone.config.api_key, mock_arg_api_key.clone()); + assert_eq!(pinecone.api_key, mock_arg_api_key.clone()); assert_eq!( - pinecone.config.controller_url, + pinecone.controller_url, mock_arg_controller_host.clone() ); - assert_eq!(pinecone.config.additional_headers, mock_arg_headers.clone()); + assert_eq!(pinecone.additional_headers, mock_arg_headers.clone()); }, ); Ok(()) } - - #[tokio::test] - async fn test_builder() { - let pinecone = PineconeClient::builder().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 = PineconeClient::builder() - .api_key("mock-api-key") - .additional_headers(HashMap::from([( - "header1".to_string(), - "value1".to_string(), - )])) - .control_plane_host("mock-controller-host") - .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 d5bf06ea3f5cfc2526a85aaa217442222f2eaab3 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Thu, 13 Jun 2024 10:45:16 -0400 Subject: [PATCH 21/28] removed parameter struct --- pinecone_sdk/src/control/create_index.rs | 103 ++++--------- pinecone_sdk/src/control/mod.rs | 3 + pinecone_sdk/src/lib.rs | 3 - .../src/models/create_index_request_params.rs | 139 ------------------ pinecone_sdk/src/models/mod.rs | 2 - pinecone_sdk/src/pinecone.rs | 2 +- 6 files changed, 35 insertions(+), 217 deletions(-) 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/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs index 5f4e1c8..354200a 100644 --- a/pinecone_sdk/src/control/create_index.rs +++ b/pinecone_sdk/src/control/create_index.rs @@ -1,69 +1,32 @@ -use crate::models::create_index_request_params::{CreateIndexParams, Spec, Metric, Cloud}; use crate::pinecone::PineconeClient; use crate::utils::errors::PineconeError; -use openapi::models::create_index_request; -use openapi::models::serverless_spec; +use openapi::models::serverless_spec::Cloud; +use openapi::models::create_index_request::Metric; use openapi::models::{CreateIndexRequest, CreateIndexRequestSpec, IndexModel, ServerlessSpec}; impl PineconeClient { - /// Creates a new index based on the provided parameters - pub async fn create_index( + /// Creates serverless index + pub async fn create_serverless_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, + name: &str, dimension: u32, - metric: Metric, - cloud: Cloud, - region: String + metric: Option, + cloud: Cloud, + region: &str ) -> Result { - // convert to openapi types - let openapi_metric: create_index_request::Metric = metric.into(); - let openapi_cloud: serverless_spec::Cloud = cloud.into(); - // create request specs let create_index_request_spec = CreateIndexRequestSpec { serverless: Some(Box::new(ServerlessSpec { - cloud: openapi_cloud, - region, + cloud, + region: region.to_string(), })), pod: None, }; let create_index_request = CreateIndexRequest { - name, + name: name.to_string(), dimension: dimension.try_into().unwrap(), - metric: Some(openapi_metric), + metric: metric, spec: Some(Box::new(create_index_request_spec)), }; @@ -119,15 +82,13 @@ mod tests { None, None, ); - let name = "index_name".to_string(); - let dimension = 10; let create_index_request = pinecone.unwrap().create_serverless_index( - name, - dimension, - Metric::Cosine, + "index_name", + 10, + Some(Metric::Cosine), Cloud::Aws, - "us-east-1".to_string() + "us-east-1" ).await; assert!(create_index_request.is_ok()); @@ -175,33 +136,31 @@ mod tests { 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; + let name = "index_name"; + let metric = Metric::Euclidean; + let cloud = Cloud::Aws; + let region = "us-east-1"; + + let result = pinecone.unwrap().create_serverless_index( + name, + 10, + Some(metric), + cloud, + region).await; match result { Ok(index) => { - assert_eq!(index.name, "index_name"); + assert_eq!(index.name, name.to_string()); 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 + Cloud::Aws ); - assert_eq!(serverless_spec.region, "us-east-1"); + assert_eq!(serverless_spec.region, region.to_string()); } Err(e) => panic!("{}", e), } diff --git a/pinecone_sdk/src/control/mod.rs b/pinecone_sdk/src/control/mod.rs index 7c509b4..3d3ae2c 100644 --- a/pinecone_sdk/src/control/mod.rs +++ b/pinecone_sdk/src/control/mod.rs @@ -3,3 +3,6 @@ pub mod create_index; /// Module for listing all indexes. pub mod list_indexes; + +pub use openapi::models::create_index_request::Metric; +pub use openapi::models::serverless_spec::Cloud; \ No newline at end of file diff --git a/pinecone_sdk/src/lib.rs b/pinecone_sdk/src/lib.rs index 79499e4..b0f956d 100644 --- a/pinecone_sdk/src/lib.rs +++ b/pinecone_sdk/src/lib.rs @@ -17,6 +17,3 @@ pub mod pinecone; /// Utility modules. pub mod utils; - -/// Structure models. -pub mod models; 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 59670dc..0000000 --- a/pinecone_sdk/src/models/create_index_request_params.rs +++ /dev/null @@ -1,139 +0,0 @@ -use openapi::models::create_index_request; -use openapi::models::serverless_spec; - -/// Request parameters for creating an index. -#[derive(Debug, PartialEq)] -pub struct CreateIndexParams { - /// The name of the index. - /// Resource name must be 1-45 characters long, start and end with an alphanumeric character, and consist only of - /// lower case alphanumeric characters or '-'. - pub name: String, - /// The dimensions of the vectors to be inserted in the index. - pub dimension: u32, - /// The distance metric to be used for similarity search. You can use 'euclidean', 'cosine', or 'dotproduct'. - pub metric: Metric, - /// The spec object defines how the index should be deployed. - pub spec: Spec, -} - -impl CreateIndexParams { - /// Create a new instance of 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, - } - } -} - -/// Spec for the index. -/// Can be either serverless or pod. -#[derive(Debug, PartialEq)] -pub enum Spec { - /// Spec for a serverless index. - Serverless { - /// The public cloud where you would like your index hosted. - cloud: Cloud, - /// The region where you would like your index to be created. - region: String, - }, - /// Spec for a pod-based index. - Pod { - /// The environment where the index is hosted. - environment: String, - /// The number of replicas. Replicas duplicate your index. They provide higher availability and throughput. - /// Replicas can be scaled up or down as your needs change. - replicas: Option, - /// The number of shards. Shards split your data across multiple pods so you can fit more data into an index. - shards: Option, - /// The type of pod to use. One of `s1`, `p1`, or `p2` appended with `.` and one of `x1`, `x2`, `x4`, or `x8`. - pod_type: String, - /// The number of pods to be used in the index. This should be equal to `shards` x `replicas`. - pods: i32, - /// Configuration for the behavior of Pinecone's internal metadata index. By default, all metadata is indexed; - /// when `metadata_config` is present, only specified metadata fields are indexed. - /// These configurations are only valid for use with pod-based indexes. - metadata_config: Option, - /// The name of the collection to be used as the source for the index. - source_collection: Option, - }, -} - -/// Distance metric to be used for similarity search. -#[derive(Debug, PartialEq)] -pub enum Metric { - /// Cosine similarity - Cosine, - /// Euclidean distance - Euclidean, - /// Dot product - Dotproduct, -} - -impl Into for Metric { - fn into(self) -> create_index_request::Metric { - match self { - Metric::Cosine => create_index_request::Metric::Cosine, - Metric::Euclidean => create_index_request::Metric::Euclidean, - Metric::Dotproduct => create_index_request::Metric::Dotproduct, - } - } -} - -/// Cloud where the index should be hosted. -#[derive(Debug, PartialEq)] -pub enum Cloud { - /// Amazon Web Services - Aws, - /// Google Cloud Platform - Gcp, - /// Microsoft Azure - Azure, -} - -impl Into for Cloud { - fn into(self) -> serverless_spec::Cloud { - match self { - Cloud::Aws => serverless_spec::Cloud::Aws, - Cloud::Gcp => serverless_spec::Cloud::Gcp, - Cloud::Azure => serverless_spec::Cloud::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(), - } - ); - } -} diff --git a/pinecone_sdk/src/models/mod.rs b/pinecone_sdk/src/models/mod.rs deleted file mode 100644 index 53a7c78..0000000 --- a/pinecone_sdk/src/models/mod.rs +++ /dev/null @@ -1,2 +0,0 @@ -/// Request parameters for create_index. -pub mod create_index_request_params; diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index 76a1283..127a8bf 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -70,7 +70,7 @@ impl PineconeClient { }, }; - let config = Config::new(api_key_option.clone(), source_tag.clone()); + let config = Config::new("".to_string(), None); let user_agent = get_user_agent(&config); Ok(PineconeClient { From e3ecb048226c66093b70bebc9ed86b411df0f8e8 Mon Sep 17 00:00:00 2001 From: Emily Yu Date: Thu, 13 Jun 2024 13:27:25 -0400 Subject: [PATCH 22/28] Use Option for cloud and use defaults --- pinecone_sdk/src/control/create_index.rs | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs index 354200a..d9d8a7c 100644 --- a/pinecone_sdk/src/control/create_index.rs +++ b/pinecone_sdk/src/control/create_index.rs @@ -11,9 +11,13 @@ impl PineconeClient { name: &str, dimension: u32, metric: Option, - cloud: Cloud, + cloud: Option, region: &str ) -> Result { + // use defaults + let metric = metric.unwrap_or(Default::default()); + let cloud = cloud.unwrap_or(Default::default()); + // create request specs let create_index_request_spec = CreateIndexRequestSpec { serverless: Some(Box::new(ServerlessSpec { @@ -26,7 +30,7 @@ impl PineconeClient { let create_index_request = CreateIndexRequest { name: name.to_string(), dimension: dimension.try_into().unwrap(), - metric: metric, + metric: Some(metric), spec: Some(Box::new(create_index_request_spec)), }; @@ -87,7 +91,7 @@ mod tests { "index_name", 10, Some(Metric::Cosine), - Cloud::Aws, + Some(Cloud::Aws), "us-east-1" ).await; assert!(create_index_request.is_ok()); @@ -146,7 +150,7 @@ mod tests { name, 10, Some(metric), - cloud, + Some(cloud), region).await; match result { From 20350174c4e1367fedd251f2e6f0bf2e6b8dd8ec Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Thu, 13 Jun 2024 13:43:51 -0400 Subject: [PATCH 23/28] remove redundant test --- pinecone_sdk/src/control/create_index.rs | 65 ------------------------ 1 file changed, 65 deletions(-) diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs index d9d8a7c..0114b5f 100644 --- a/pinecone_sdk/src/control/create_index.rs +++ b/pinecone_sdk/src/control/create_index.rs @@ -49,7 +49,6 @@ impl PineconeClient { #[cfg(test)] mod tests { use super::*; - use core::panic; use mockito::mock; use tokio; @@ -105,68 +104,4 @@ mod tests { 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 = PineconeClient::new( - Some("api_key".to_string()), - Some(mockito::server_url()), - None, - None, - ); - - let name = "index_name"; - let metric = Metric::Euclidean; - let cloud = Cloud::Aws; - let region = "us-east-1"; - - let result = pinecone.unwrap().create_serverless_index( - name, - 10, - Some(metric), - Some(cloud), - region).await; - - match result { - Ok(index) => { - assert_eq!(index.name, name.to_string()); - 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, - Cloud::Aws - ); - assert_eq!(serverless_spec.region, region.to_string()); - } - Err(e) => panic!("{}", e), - } - } } From 34871aff9eaadc24dbd28345f65c20cc3f2c0e10 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Thu, 13 Jun 2024 15:52:07 -0400 Subject: [PATCH 24/28] remove Option for Cloud and Metric --- pinecone_sdk/src/control/create_index.rs | 69 ++++++++++++++++++++---- 1 file changed, 59 insertions(+), 10 deletions(-) diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs index 0114b5f..9bbdd1c 100644 --- a/pinecone_sdk/src/control/create_index.rs +++ b/pinecone_sdk/src/control/create_index.rs @@ -10,14 +10,10 @@ impl PineconeClient { &self, name: &str, dimension: u32, - metric: Option, - cloud: Option, + metric: Metric, + cloud: Cloud, region: &str ) -> Result { - // use defaults - let metric = metric.unwrap_or(Default::default()); - let cloud = cloud.unwrap_or(Default::default()); - // create request specs let create_index_request_spec = CreateIndexRequestSpec { serverless: Some(Box::new(ServerlessSpec { @@ -62,7 +58,7 @@ mod tests { { "name": "index_name", "dimension": 10, - "metric": "cosine", + "metric": "euclidean", "host": "host1", "spec": { "serverless": { @@ -89,8 +85,8 @@ mod tests { let create_index_request = pinecone.unwrap().create_serverless_index( "index_name", 10, - Some(Metric::Cosine), - Some(Cloud::Aws), + Metric::Cosine, + Cloud::Aws, "us-east-1" ).await; assert!(create_index_request.is_ok()); @@ -98,10 +94,63 @@ mod tests { 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); + assert_eq!(create_index_req.metric, openapi::models::index_model::Metric::Euclidean); 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_defaults() { + 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": "gcp", + "region": "us-east-1" + } + }, + "status": { + "ready": true, + "state": "Initializing" + } + } + "#, + ) + .create(); + + let pinecone = PineconeClient::new( + Some("api_key".to_string()), + Some(mockito::server_url()), + None, + None, + ); + + let create_index_request = pinecone.unwrap().create_serverless_index( + "index_name", + 10, + Default::default(), + Default::default(), + "us-east-1" + ).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::Gcp); + assert_eq!(spec.region, "us-east-1"); + } } From cadf7b2c4b73c76104e3fcc0fec68d0108705f2e Mon Sep 17 00:00:00 2001 From: Emily Yu Date: Thu, 13 Jun 2024 16:40:43 -0400 Subject: [PATCH 25/28] Condense control code --- pinecone_sdk/src/control.rs | 273 +++++++++++++++++++++++ pinecone_sdk/src/control/create_index.rs | 156 ------------- pinecone_sdk/src/control/list_indexes.rs | 124 ---------- pinecone_sdk/src/control/mod.rs | 8 - pinecone_sdk/tests/integration_test.rs | 2 +- 5 files changed, 274 insertions(+), 289 deletions(-) create mode 100644 pinecone_sdk/src/control.rs delete mode 100644 pinecone_sdk/src/control/create_index.rs delete mode 100644 pinecone_sdk/src/control/list_indexes.rs delete mode 100644 pinecone_sdk/src/control/mod.rs diff --git a/pinecone_sdk/src/control.rs b/pinecone_sdk/src/control.rs new file mode 100644 index 0000000..b961dbe --- /dev/null +++ b/pinecone_sdk/src/control.rs @@ -0,0 +1,273 @@ +use crate::pinecone::PineconeClient; +use crate::utils::errors::PineconeError; +use openapi::apis::manage_indexes_api; +use openapi::apis::manage_indexes_api::ListIndexesError; +use openapi::apis::Error; +use openapi::models; +use openapi::models::{CreateIndexRequest, CreateIndexRequestSpec, IndexModel, ServerlessSpec}; + +pub use openapi::models::create_index_request::Metric; +pub use openapi::models::serverless_spec::Cloud; + +impl PineconeClient { + /// Creates serverless index + pub async fn create_serverless_index( + &self, + name: &str, + dimension: u32, + metric: Metric, + cloud: Cloud, + region: &str, + ) -> Result { + // create request specs + let create_index_request_spec = CreateIndexRequestSpec { + serverless: Some(Box::new(ServerlessSpec { + cloud, + region: region.to_string(), + })), + pod: None, + }; + + let create_index_request = CreateIndexRequest { + name: name.to_string(), + dimension: dimension.try_into().unwrap(), + metric: Some(metric), + 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 }), + } + } + /// Lists all indexes. + /// + /// The results include a description of all indexes in your project, including the + /// index name, dimension, metric, status, and spec. + /// + /// :return: Returns an `IndexList` object, which is iterable and contains a + /// list of `IndexDescription` objects. It also has a convenience method `names()` + /// which returns a list of index names. + /// + /// ### Example + /// + /// ``` + /// # use pinecone_sdk::pinecone::PineconeClient; + /// # use pinecone_sdk::utils::errors::PineconeError; + /// # #[tokio::main] + /// # async fn main() -> Result<(), PineconeError>{ + /// # // Create a Pinecone client with the API key and controller host. + /// # let pinecone = PineconeClient::new(None, None, None, None).unwrap(); + /// # + /// // List all indexes in the project. + /// let index_list = pinecone.list_indexes(); + /// # Ok(()) + /// # } + /// ``` + + pub async fn list_indexes(&self) -> Result> { + let response = manage_indexes_api::list_indexes(&self.openapi_config()).await?; + println!("{:?}", response); + Ok(response) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use mockito::mock; + use models::IndexList; + use tokio; + + #[tokio::test] + 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": "euclidean", + "host": "host1", + "spec": { + "serverless": { + "cloud": "aws", + "region": "us-east-1" + } + }, + "status": { + "ready": true, + "state": "Initializing" + } + } + "#, + ) + .create(); + + let pinecone = PineconeClient::new( + Some("api_key".to_string()), + Some(mockito::server_url()), + None, + None, + ); + + let create_index_request = pinecone + .unwrap() + .create_serverless_index("index_name", 10, Metric::Cosine, Cloud::Aws, "us-east-1") + .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::Euclidean + ); + + 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_defaults() { + 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": "gcp", + "region": "us-east-1" + } + }, + "status": { + "ready": true, + "state": "Initializing" + } + } + "#, + ) + .create(); + + let pinecone = PineconeClient::new( + Some("api_key".to_string()), + Some(mockito::server_url()), + None, + None, + ); + + let create_index_request = pinecone + .unwrap() + .create_serverless_index( + "index_name", + 10, + Default::default(), + Default::default(), + "us-east-1", + ) + .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::Gcp); + assert_eq!(spec.region, "us-east-1"); + } + + #[tokio::test] + async fn test_list_indexes() -> Result<(), PineconeError> { + // Create a mock server + let _m = mock("GET", "/indexes") + .with_status(200) + .with_header("content-type", "application/json") + .with_body( + r#" + { + "indexes": [ + { + "name": "index1", + "dimension": 1536, + "metric": "cosine", + "host": "host1", + "spec": {}, + "status": { + "ready": false, + "state": "Initializing" + } + }, + { + "name": "index2", + "dimension": 1536, + "metric": "cosine", + "host": "host2", + "spec": {}, + "status": { + "ready": false, + "state": "Initializing" + } + } + ] + } + "#, + ) + .create(); + + // Construct Pinecone instance with the mock server URL + let api_key = "test_api_key".to_string(); + let pinecone = PineconeClient::new(Some(api_key), Some(mockito::server_url()), None, None) + .expect("Failed to create Pinecone instance"); + + // Call list_indexes and verify the result + let index_list = pinecone + .list_indexes() + .await + .expect("Failed to list indexes"); + + let expected = IndexList { + // name: String, dimension: i32, metric: Metric, host: String, spec: models::IndexModelSpec, status: models::IndexModelStatus) + indexes: Some(vec![ + IndexModel::new( + "index1".to_string(), + 1536, + openapi::models::index_model::Metric::Cosine, + "host1".to_string(), + models::IndexModelSpec::default(), + models::IndexModelStatus::default(), + ), + IndexModel::new( + "index2".to_string(), + 1536, + openapi::models::index_model::Metric::Cosine, + "host2".to_string(), + models::IndexModelSpec::default(), + models::IndexModelStatus::default(), + ), + ]), + }; + assert_eq!(index_list, expected); + + Ok(()) + } +} diff --git a/pinecone_sdk/src/control/create_index.rs b/pinecone_sdk/src/control/create_index.rs deleted file mode 100644 index 9bbdd1c..0000000 --- a/pinecone_sdk/src/control/create_index.rs +++ /dev/null @@ -1,156 +0,0 @@ -use crate::pinecone::PineconeClient; -use crate::utils::errors::PineconeError; -use openapi::models::serverless_spec::Cloud; -use openapi::models::create_index_request::Metric; -use openapi::models::{CreateIndexRequest, CreateIndexRequestSpec, IndexModel, ServerlessSpec}; - -impl PineconeClient { - /// Creates serverless index - pub async fn create_serverless_index( - &self, - name: &str, - dimension: u32, - metric: Metric, - cloud: Cloud, - region: &str - ) -> Result { - // create request specs - let create_index_request_spec = CreateIndexRequestSpec { - serverless: Some(Box::new(ServerlessSpec { - cloud, - region: region.to_string(), - })), - pod: None, - }; - - let create_index_request = CreateIndexRequest { - name: name.to_string(), - dimension: dimension.try_into().unwrap(), - metric: Some(metric), - 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 mockito::mock; - use tokio; - - #[tokio::test] - 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": "euclidean", - "host": "host1", - "spec": { - "serverless": { - "cloud": "aws", - "region": "us-east-1" - } - }, - "status": { - "ready": true, - "state": "Initializing" - } - } - "#, - ) - .create(); - - let pinecone = PineconeClient::new( - Some("api_key".to_string()), - Some(mockito::server_url()), - None, - None, - ); - - let create_index_request = pinecone.unwrap().create_serverless_index( - "index_name", - 10, - Metric::Cosine, - Cloud::Aws, - "us-east-1" - ).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::Euclidean); - - 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_defaults() { - 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": "gcp", - "region": "us-east-1" - } - }, - "status": { - "ready": true, - "state": "Initializing" - } - } - "#, - ) - .create(); - - let pinecone = PineconeClient::new( - Some("api_key".to_string()), - Some(mockito::server_url()), - None, - None, - ); - - let create_index_request = pinecone.unwrap().create_serverless_index( - "index_name", - 10, - Default::default(), - Default::default(), - "us-east-1" - ).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::Gcp); - assert_eq!(spec.region, "us-east-1"); - } -} diff --git a/pinecone_sdk/src/control/list_indexes.rs b/pinecone_sdk/src/control/list_indexes.rs deleted file mode 100644 index 4b465d8..0000000 --- a/pinecone_sdk/src/control/list_indexes.rs +++ /dev/null @@ -1,124 +0,0 @@ -use crate::pinecone::PineconeClient; -use openapi::apis::manage_indexes_api; -use openapi::apis::manage_indexes_api::ListIndexesError; -use openapi::apis::Error; -use openapi::models; - -impl PineconeClient { - /// Lists all indexes. - /// - /// The results include a description of all indexes in your project, including the - /// index name, dimension, metric, status, and spec. - /// - /// :return: Returns an `IndexList` object, which is iterable and contains a - /// list of `IndexDescription` objects. It also has a convenience method `names()` - /// which returns a list of index names. - /// - /// ### Example - /// - /// ``` - /// # use pinecone_sdk::pinecone::PineconeClient; - /// # use pinecone_sdk::utils::errors::PineconeError; - /// # #[tokio::main] - /// # async fn main() -> Result<(), PineconeError>{ - /// # // Create a Pinecone client with the API key and controller host. - /// # let pinecone = PineconeClient::new(None, None, None, None).unwrap(); - /// # - /// // List all indexes in the project. - /// let index_list = pinecone.list_indexes(); - /// # Ok(()) - /// # } - /// ``` - - pub async fn list_indexes(&self) -> Result> { - let response = manage_indexes_api::list_indexes(&self.openapi_config()).await?; - println!("{:?}", response); - Ok(response) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::control::list_indexes::models::index_model::Metric; - use crate::utils::errors::PineconeError; - use mockito::mock; - use openapi::models::IndexList; - use openapi::models::IndexModel; - use tokio; - - #[tokio::test] - async fn test_list_indexes() -> Result<(), PineconeError> { - // Create a mock server - let _m = mock("GET", "/indexes") - .with_status(200) - .with_header("content-type", "application/json") - .with_body( - r#" - { - "indexes": [ - { - "name": "index1", - "dimension": 1536, - "metric": "cosine", - "host": "host1", - "spec": {}, - "status": { - "ready": false, - "state": "Initializing" - } - }, - { - "name": "index2", - "dimension": 1536, - "metric": "cosine", - "host": "host2", - "spec": {}, - "status": { - "ready": false, - "state": "Initializing" - } - } - ] - } - "#, - ) - .create(); - - // Construct Pinecone instance with the mock server URL - let api_key = "test_api_key".to_string(); - let pinecone = PineconeClient::new(Some(api_key), Some(mockito::server_url()), None, None) - .expect("Failed to create Pinecone instance"); - - // Call list_indexes and verify the result - let index_list = pinecone - .list_indexes() - .await - .expect("Failed to list indexes"); - - let expected = IndexList { - // name: String, dimension: i32, metric: Metric, host: String, spec: models::IndexModelSpec, status: models::IndexModelStatus) - indexes: Some(vec![ - IndexModel::new( - "index1".to_string(), - 1536, - Metric::Cosine, - "host1".to_string(), - models::IndexModelSpec::default(), - models::IndexModelStatus::default(), - ), - IndexModel::new( - "index2".to_string(), - 1536, - Metric::Cosine, - "host2".to_string(), - models::IndexModelSpec::default(), - models::IndexModelStatus::default(), - ), - ]), - }; - assert_eq!(index_list, expected); - - Ok(()) - } -} diff --git a/pinecone_sdk/src/control/mod.rs b/pinecone_sdk/src/control/mod.rs deleted file mode 100644 index 3d3ae2c..0000000 --- a/pinecone_sdk/src/control/mod.rs +++ /dev/null @@ -1,8 +0,0 @@ -/// Module for creating an index. -pub mod create_index; - -/// Module for listing all indexes. -pub mod list_indexes; - -pub use openapi::models::create_index_request::Metric; -pub use openapi::models::serverless_spec::Cloud; \ No newline at end of file diff --git a/pinecone_sdk/tests/integration_test.rs b/pinecone_sdk/tests/integration_test.rs index 6fe9b1c..de46494 100644 --- a/pinecone_sdk/tests/integration_test.rs +++ b/pinecone_sdk/tests/integration_test.rs @@ -2,7 +2,7 @@ use pinecone_sdk::pinecone::PineconeClient; use pinecone_sdk::utils::errors::PineconeError; #[tokio::test] -async fn test_list_index_env() -> Result<(), PineconeError> { +async fn test_list_index() -> Result<(), PineconeError> { let pinecone = PineconeClient::new(None, None, None, None).unwrap(); let _ = pinecone .list_indexes() From 24619e6d3863340bb4f70ffa97d3d0012edc6523 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Thu, 13 Jun 2024 17:12:32 -0400 Subject: [PATCH 26/28] pass api key to config --- pinecone_sdk/src/pinecone.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index 127a8bf..025d32f 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -70,7 +70,7 @@ impl PineconeClient { }, }; - let config = Config::new("".to_string(), None); + let config = Config::new(api_key_option.clone(), None); let user_agent = get_user_agent(&config); Ok(PineconeClient { From 90e325424cd805665b39480f2ca546348a11e3d7 Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Thu, 13 Jun 2024 17:14:09 -0400 Subject: [PATCH 27/28] changed api_key_opt to api_key_str --- pinecone_sdk/src/pinecone.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index 025d32f..73e3ddd 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -42,7 +42,7 @@ impl PineconeClient { source_tag: Option, ) -> Result { // get api key - let api_key_option = match api_key { + let api_key_str = match api_key { Some(key) => key, None => match std::env::var("PINECONE_API_KEY") { Ok(key) => key, @@ -70,11 +70,11 @@ impl PineconeClient { }, }; - let config = Config::new(api_key_option.clone(), None); + let config = Config::new(api_key_str.clone(), None); let user_agent = get_user_agent(&config); Ok(PineconeClient { - api_key: api_key_option, + api_key: api_key_str, controller_url: controller_host, additional_headers, source_tag, From e0accd49f9d03a13ce1e2a60305288c87bd8ee6b Mon Sep 17 00:00:00 2001 From: Erica Wang Date: Fri, 14 Jun 2024 09:53:20 -0400 Subject: [PATCH 28/28] fixed config construction in PineconeClient --- pinecone_sdk/src/pinecone.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index 73e3ddd..4232971 100644 --- a/pinecone_sdk/src/pinecone.rs +++ b/pinecone_sdk/src/pinecone.rs @@ -70,7 +70,12 @@ impl PineconeClient { }, }; - let config = Config::new(api_key_str.clone(), None); + let config = Config { + api_key: api_key_str.clone(), + controller_url: controller_host.clone(), + additional_headers: additional_headers.clone(), + source_tag: source_tag.clone(), + }; let user_agent = get_user_agent(&config); Ok(PineconeClient {