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/list_indexes.rs b/pinecone_sdk/src/control/list_indexes.rs deleted file mode 100644 index d24f76e..0000000 --- a/pinecone_sdk/src/control/list_indexes.rs +++ /dev/null @@ -1,124 +0,0 @@ -use crate::pinecone::Pinecone; -use openapi::apis::manage_indexes_api; -use openapi::apis::manage_indexes_api::ListIndexesError; -use openapi::apis::Error; -use openapi::models; - -impl Pinecone { - /// 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::Pinecone; - /// # 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(); - /// # - /// // 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 = Pinecone::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 bc5d416..0000000 --- a/pinecone_sdk/src/control/mod.rs +++ /dev/null @@ -1 +0,0 @@ -mod list_indexes; diff --git a/pinecone_sdk/src/pinecone.rs b/pinecone_sdk/src/pinecone.rs index 670ed2f..4232971 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 Pinecone { - /// Configuration for the Pinecone SDK struct. - config: Config, - - /// OpenAPI configuration object. - openapi_config: Configuration, +pub struct PineconeClient { + api_key: String, + controller_url: String, + additional_headers: HashMap, + source_tag: Option, + user_agent: Option, } -impl Pinecone { - /// The `Pinecone` struct is the main entry point for interacting with Pinecone via this Rust SDK. +impl PineconeClient { + /// 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 @@ -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, @@ -42,14 +42,14 @@ impl Pinecone { source_tag: Option, ) -> Result { // get api key - let api_key = match api_key { + let api_key_str = 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( @@ -71,33 +71,33 @@ impl Pinecone { }; let config = Config { - api_key: api_key.clone(), + api_key: api_key_str.clone(), controller_url: controller_host.clone(), - additional_headers, - source_tag, + additional_headers: additional_headers.clone(), + source_tag: source_tag.clone(), }; - let user_agent = get_user_agent(&config); - - let openapi_config = Configuration { - base_path: controller_host, + + Ok(PineconeClient { + api_key: api_key_str, + controller_url: controller_host, + additional_headers, + source_tag, user_agent: Some(user_agent), - api_key: Some(ApiKey { - prefix: None, - key: api_key, - }), - ..Default::default() - }; - - Ok(Pinecone { - config, - openapi_config, }) } /// Returns the OpenAPI configuration object. - pub fn openapi_config(&self) -> &Configuration { - &self.openapi_config + 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() + } } } @@ -111,7 +111,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()), @@ -119,7 +119,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(()) } @@ -130,7 +134,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()), @@ -138,7 +142,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(()) @@ -149,7 +157,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()), @@ -167,7 +175,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()), @@ -175,7 +183,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(()) } @@ -190,10 +198,10 @@ 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()); + assert_eq!(pinecone.controller_url, mock_controller_host); }, ); @@ -205,7 +213,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()), @@ -216,7 +224,7 @@ mod tests { ); assert_eq!( - pinecone.config.controller_url, + pinecone.controller_url, "https://api.pinecone.io".to_string() ); }); @@ -233,7 +241,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()), @@ -241,7 +249,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(()) } @@ -259,7 +267,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, @@ -267,7 +275,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); }, ); @@ -280,7 +288,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, @@ -303,7 +311,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()), @@ -311,7 +319,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(()) @@ -345,7 +353,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()), @@ -353,12 +361,12 @@ 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()); }, ); diff --git a/pinecone_sdk/src/utils/errors.rs b/pinecone_sdk/src/utils/errors.rs index 03a67e7..08ce7a0 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::*; /// PineconeError is the error type for all Pinecone SDK errors. @@ -7,10 +8,69 @@ 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: Cloud name. + cloud: String, + }, + /// InvalidHeadersError: Provided headers are not valid. Expects JSON. #[snafu(display("Failed to parse headers: {}", json_error))] InvalidHeadersError { /// json_error: Error object for JSON parsing error. json_error: serde_json::Error, }, + + /// InvalidMetricError: Provided metric is not valid. + #[snafu(display("Invalid metric '{}'.", metric))] + 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, +} + +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, + _ => false, + } + } } diff --git a/pinecone_sdk/tests/integration_test.rs b/pinecone_sdk/tests/integration_test.rs index e539b81..de46494 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(); +async fn test_list_index() -> Result<(), PineconeError> { + let pinecone = PineconeClient::new(None, None, None, None).unwrap(); let _ = pinecone .list_indexes() .await