diff --git a/Cargo.lock b/Cargo.lock index 4f2b2ae..ea8078e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1615,14 +1615,13 @@ dependencies = [ "serde_json", "serde_with", "serial_test", - "snafu", + "strum", "temp-env", "thiserror", "tokio", "tonic", "tonic-build", "url", - "uuid", ] [[package]] @@ -2228,27 +2227,6 @@ version = "1.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" -[[package]] -name = "snafu" -version = "0.8.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b835cb902660db3415a672d862905e791e54d306c6e8189168c7f3d9ae1c79d" -dependencies = [ - "snafu-derive", -] - -[[package]] -name = "snafu-derive" -version = "0.8.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38d1e02fca405f6280643174a50c942219f0bbf4dbf7d480f1dd864d6f211ae5" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "syn 2.0.71", -] - [[package]] name = "socket2" version = "0.4.10" @@ -2294,6 +2272,28 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "strum" +version = "0.27.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f64def088c51c9510a8579e3c5d67c65349dcf755e5479ad3d010aa6454e2c32" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.27.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c77a8c5abcaf0f9ce05d62342b7d298c346515365c36b673df4ebe3ced01fde8" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "rustversion", + "syn 2.0.71", +] + [[package]] name = "subtle" version = "2.6.1" @@ -2724,16 +2724,6 @@ dependencies = [ "percent-encoding", ] -[[package]] -name = "uuid" -version = "1.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81dfa00651efa65069b0b6b651f4aaa31ba9e3c3ce0137aaad053604ee7e0314" -dependencies = [ - "getrandom", - "serde", -] - [[package]] name = "value-bag" version = "1.9.0" diff --git a/Cargo.toml b/Cargo.toml index 168a01b..74da3a5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,20 +18,15 @@ exclude = ["tests/*"] tokio = { version = "1", features = ["full"] } regex = "1.10" serde_json = "1.0" -snafu = "0.8" rand = "0.8" tonic = { version = "0.11", features = ["tls", "transport", "tls-roots"] } prost = "0.12" prost-types = "0.12" -# reqwest = "0.12" once_cell = "1.19" - -# openapi serde = { version = "^1.0", features = ["derive"] } serde_with = { version = "3.12", features = ["base64"] } -# serde_json = "^1.0" +strum = { version = "0.27", features = ["derive"] } url = "^2.5" -uuid = { version = "^1.8", features = ["serde", "v4"] } reqwest = { version = "^0.12", features = ["json", "multipart"] } thiserror = "1.0.63" anyhow = "1.0.86" diff --git a/src/lib.rs b/src/lib.rs index 1fa337e..b67d42f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -48,9 +48,9 @@ //! //! ```no_run //! use pinecone_sdk::pinecone; -//! use pinecone_sdk::models::{Cloud, DeletionProtection, IndexModel, Metric, WaitPolicy}; +//! use pinecone_sdk::models::{Cloud, DeletionProtection, IndexModel, Metric, VectorType, WaitPolicy}; //! use pinecone_sdk::utils::errors::PineconeError; -//! # async fn create_index_and_collection() -> Result<(), PineconeError> { +//! # async fn create_index() -> Result<(), PineconeError> { //! let client: pinecone::PineconeClient = //! pinecone::default_client().expect("Failed to create PineconeClient"); //! @@ -63,17 +63,15 @@ //! "us-east-1", //! DeletionProtection::Disabled, //! WaitPolicy::NoWait, +//! VectorType::Dense, +//! None, //! ) //! .await?; //! -//! let collection = client.create_collection("my-collection-name", "my-previous-index-name").await?; -//! //! let index_description = client.describe_index("index-name").await?; -//! let collection_description = client.describe_collection("my-collection-name").await?; //! let indexes = client.list_indexes().await?; //! //! println!("Index description: {:?}", index_description); -//! println!("Collection description: {:?}", collection_description); //! println!("Index list: {:?}", indexes); //! //! # Ok(()) diff --git a/src/models/cloud.rs b/src/models/cloud.rs new file mode 100644 index 0000000..85ed993 --- /dev/null +++ b/src/models/cloud.rs @@ -0,0 +1,58 @@ +use serde::{Deserialize, Serialize}; +use strum::{Display, EnumString}; + +/// The public cloud where you would like your index hosted. +#[derive( + Debug, Default, Clone, Copy, PartialEq, Eq, Display, EnumString, Serialize, Deserialize, +)] +#[serde(rename_all = "lowercase")] +#[strum(serialize_all = "lowercase")] +pub enum Cloud { + /// GCP + #[default] + Gcp, + /// AWS + Aws, + /// Azure + Azure, +} + +impl From for Cloud { + fn from(cloud: crate::openapi::models::serverless_spec::Cloud) -> Self { + match cloud { + crate::openapi::models::serverless_spec::Cloud::Gcp => Cloud::Gcp, + crate::openapi::models::serverless_spec::Cloud::Aws => Cloud::Aws, + crate::openapi::models::serverless_spec::Cloud::Azure => Cloud::Azure, + } + } +} + +impl From for crate::openapi::models::serverless_spec::Cloud { + fn from(cloud: Cloud) -> Self { + match cloud { + Cloud::Gcp => crate::openapi::models::serverless_spec::Cloud::Gcp, + Cloud::Aws => crate::openapi::models::serverless_spec::Cloud::Aws, + Cloud::Azure => crate::openapi::models::serverless_spec::Cloud::Azure, + } + } +} + +impl From for Cloud { + fn from(cloud: crate::openapi::models::create_index_for_model_request::Cloud) -> Self { + match cloud { + crate::openapi::models::create_index_for_model_request::Cloud::Gcp => Cloud::Gcp, + crate::openapi::models::create_index_for_model_request::Cloud::Aws => Cloud::Aws, + crate::openapi::models::create_index_for_model_request::Cloud::Azure => Cloud::Azure, + } + } +} + +impl From for crate::openapi::models::create_index_for_model_request::Cloud { + fn from(cloud: Cloud) -> Self { + match cloud { + Cloud::Gcp => crate::openapi::models::create_index_for_model_request::Cloud::Gcp, + Cloud::Aws => crate::openapi::models::create_index_for_model_request::Cloud::Aws, + Cloud::Azure => crate::openapi::models::create_index_for_model_request::Cloud::Azure, + } + } +} diff --git a/src/models/index_model.rs b/src/models/index_model.rs index 9504baf..b441bce 100644 --- a/src/models/index_model.rs +++ b/src/models/index_model.rs @@ -1,5 +1,8 @@ -use super::{DeletionProtection, IndexModelSpec, IndexModelStatus, Metric}; +use super::{DeletionProtection, IndexModelSpec, IndexModelStatus, Metric, VectorType}; use crate::openapi::models::index_model::IndexModel as OpenApiIndexModel; +use crate::openapi::models::{CreateIndexForModelRequestEmbed, ModelIndexEmbed}; +use serde_with::serde_derive::Serialize; +use std::collections::HashMap; /// IndexModel : The IndexModel describes the configuration and status of a Pinecone index. #[derive(Clone, Default, Debug, PartialEq)] @@ -18,6 +21,12 @@ pub struct IndexModel { pub spec: IndexModelSpec, /// Index model specs pub status: IndexModelStatus, + /// Index tags + pub tags: Option>, + /// Index embedding configuration + pub embed: Option, + /// Index vector type + pub vector_type: VectorType, } impl From for IndexModel { @@ -30,6 +39,84 @@ impl From for IndexModel { deletion_protection: openapi_index_model.deletion_protection, spec: *openapi_index_model.spec, status: *openapi_index_model.status, + tags: openapi_index_model.tags, + embed: openapi_index_model.embed.map(|emb| *emb), + vector_type: openapi_index_model.vector_type, + } + } +} + +/// A field mapping entry by type. +#[derive(Clone, Debug)] +pub enum FieldMapEntry { + /// The name of the text field from your document model that is embedded. + TextField(String), +} + +/// A model parameter value of a specific type. +#[derive(Clone, Debug, Serialize)] +pub enum ModelParameterValue { + /// A string value type + StringVal(String), + /// An integer value type + IntVal(i32), + /// A floating point value type + FloatVal(f32), + /// A boolean value type. + BoolVal(bool), +} + +/// Configuration options for the index with integrated embedding. +#[derive(Clone, Debug)] +pub struct CreateIndexForModelOptions { + /// The name of the embedding model to use for the index. + pub model: String, + /// Identifies the name of the field from your document model that will be embedded. (Only one + /// field is supported for now.) + pub field_map: Vec, + /// The distance metric to be used for similarity search. You can use 'euclidean', 'cosine', or 'dotproduct'. If not specified, the metric will be defaulted according to the model. Cannot be updated once set. + pub metric: Option, + /// The desired vector dimension, if supported by the model. + pub dimension: Option, + /// The read parameters for the embedding model. + pub read_parameters: Option>, + /// The write parameters for the embedding model. + pub write_parameters: Option>, +} + +impl From for CreateIndexForModelRequestEmbed { + fn from(options: CreateIndexForModelOptions) -> Self { + let field_map = options + .field_map + .into_iter() + .map(|entry| match entry { + FieldMapEntry::TextField(field_name) => { + ("text", serde_json::Value::String(field_name)) + } + }) + .collect(); + + let read_parameters = options.read_parameters.map(|params| { + params + .into_iter() + .map(|(key, value)| (key, serde_json::to_value(value).unwrap())) + .collect() + }); + + let write_parameters = options.write_parameters.map(|params| { + params + .into_iter() + .map(|(key, value)| (key, serde_json::to_value(value).unwrap())) + .collect() + }); + + CreateIndexForModelRequestEmbed { + model: options.model, + field_map, + metric: options.metric.map(|m| m.into()), + read_parameters, + write_parameters, + dimension: options.dimension, } } } diff --git a/src/models/metric.rs b/src/models/metric.rs index 4fcabdc..289e197 100644 --- a/src/models/metric.rs +++ b/src/models/metric.rs @@ -52,3 +52,35 @@ impl From for ResponseMetric { } } } + +impl From for crate::openapi::models::create_index_for_model_request_embed::Metric { + fn from(model: Metric) -> Self { + match model { + Metric::Cosine => { + crate::openapi::models::create_index_for_model_request_embed::Metric::Cosine + } + Metric::Euclidean => { + crate::openapi::models::create_index_for_model_request_embed::Metric::Euclidean + } + Metric::Dotproduct => { + crate::openapi::models::create_index_for_model_request_embed::Metric::Dotproduct + } + } + } +} + +impl From for Metric { + fn from(model: crate::openapi::models::create_index_for_model_request_embed::Metric) -> Self { + match model { + crate::openapi::models::create_index_for_model_request_embed::Metric::Cosine => { + Metric::Cosine + } + crate::openapi::models::create_index_for_model_request_embed::Metric::Euclidean => { + Metric::Euclidean + } + crate::openapi::models::create_index_for_model_request_embed::Metric::Dotproduct => { + Metric::Dotproduct + } + } + } +} diff --git a/src/models/mod.rs b/src/models/mod.rs index c05b564..41fed44 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -11,7 +11,9 @@ mod namespace; pub use self::namespace::Namespace; mod index_model; -pub use self::index_model::IndexModel; +pub use self::index_model::{ + CreateIndexForModelOptions, FieldMapEntry, IndexModel, ModelParameterValue, +}; mod index_list; pub use self::index_list::IndexList; @@ -22,11 +24,19 @@ pub use self::wait_policy::WaitPolicy; mod embedding; pub use self::embedding::Embedding; +mod vector_type; +pub use self::vector_type::VectorType; + +mod cloud; +pub use self::cloud::Cloud; + pub use crate::openapi::models::{ - index_model_status::State, serverless_spec::Cloud, CollectionList, CollectionModel, - ConfigureIndexRequest, ConfigureIndexRequestSpec, ConfigureIndexRequestSpecPod, - CreateCollectionRequest, DeletionProtection, EmbedRequestParameters, IndexModelSpec, - IndexModelStatus, IndexSpec, PodSpec, PodSpecMetadataConfig, ServerlessSpec, + index_model_status::State, CollectionList, CollectionModel, ConfigureIndexRequest, + ConfigureIndexRequestSpec, ConfigureIndexRequestSpecPod, CreateCollectionRequest, + DeletionProtection, EmbedRequestParameters, IndexModelSpec, IndexModelStatus, IndexSpec, + PodSpec, PodSpecMetadataConfig, SearchRecordsRequest, SearchRecordsRequestQuery, + SearchRecordsRequestRerank, SearchRecordsResponse, ServerlessSpec, UpsertRecord, + UpsertResponse as UpsertRecordResponse, }; pub use crate::protos::{ diff --git a/src/models/vector_type.rs b/src/models/vector_type.rs new file mode 100644 index 0000000..a20be6c --- /dev/null +++ b/src/models/vector_type.rs @@ -0,0 +1,16 @@ +use serde::{Deserialize, Serialize}; +use strum::{Display, EnumString}; + +/// Vector type, either dense or sparse. +#[derive( + Debug, Default, Clone, Copy, PartialEq, Eq, Display, EnumString, Serialize, Deserialize, +)] +#[serde(rename_all = "lowercase")] +#[strum(serialize_all = "lowercase")] +pub enum VectorType { + /// Dense vector type + #[default] + Dense, + /// Sparse vector type + Sparse, +} diff --git a/src/openapi/models/create_index_request.rs b/src/openapi/models/create_index_request.rs index 2e1044a..c33a0d0 100644 --- a/src/openapi/models/create_index_request.rs +++ b/src/openapi/models/create_index_request.rs @@ -8,6 +8,7 @@ * Generated by: https://openapi-generator.tech */ +use crate::models::VectorType; use crate::openapi::models; use serde::{Deserialize, Serialize}; @@ -35,7 +36,7 @@ pub struct CreateIndexRequest { pub spec: Option>, /// The index vector type. You can use 'dense' or 'sparse'. If 'dense', the vector dimension must be specified. If 'sparse', the vector dimension should not be specified. #[serde(rename = "vector_type", skip_serializing_if = "Option::is_none")] - pub vector_type: Option, + pub vector_type: Option, } impl CreateIndexRequest { diff --git a/src/openapi/models/index_description.rs b/src/openapi/models/index_description.rs index ad18a39..51e26ad 100644 --- a/src/openapi/models/index_description.rs +++ b/src/openapi/models/index_description.rs @@ -8,6 +8,7 @@ * Generated by: https://openapi-generator.tech */ +use crate::models::VectorType; use crate::openapi::models; use serde::{Deserialize, Serialize}; @@ -31,7 +32,7 @@ pub struct IndexDescription { pub metric: Option, /// The type of vectors stored in the index. #[serde(rename = "vectorType", skip_serializing_if = "Option::is_none")] - pub vector_type: Option, + pub vector_type: Option, } impl IndexDescription { diff --git a/src/openapi/models/index_model.rs b/src/openapi/models/index_model.rs index 7cfc591..497159a 100644 --- a/src/openapi/models/index_model.rs +++ b/src/openapi/models/index_model.rs @@ -8,6 +8,7 @@ * Generated by: https://openapi-generator.tech */ +use crate::models::VectorType; use crate::openapi::models; use serde::{Deserialize, Serialize}; @@ -42,7 +43,7 @@ pub struct IndexModel { pub status: Box, /// The index vector type. You can use 'dense' or 'sparse'. If 'dense', the vector dimension must be specified. If 'sparse', the vector dimension should not be specified. #[serde(rename = "vector_type")] - pub vector_type: String, + pub vector_type: VectorType, } impl IndexModel { @@ -53,7 +54,7 @@ impl IndexModel { host: String, spec: models::IndexModelSpec, status: models::IndexModelStatus, - vector_type: String, + vector_type: VectorType, ) -> IndexModel { IndexModel { name, diff --git a/src/openapi/models/model_index_embed.rs b/src/openapi/models/model_index_embed.rs index ea3b2b9..d045c8e 100644 --- a/src/openapi/models/model_index_embed.rs +++ b/src/openapi/models/model_index_embed.rs @@ -8,6 +8,7 @@ * Generated by: https://openapi-generator.tech */ +use crate::models::VectorType; use crate::openapi::models; use serde::{Deserialize, Serialize}; @@ -25,7 +26,7 @@ pub struct ModelIndexEmbed { pub dimension: Option, /// The index vector type. You can use 'dense' or 'sparse'. If 'dense', the vector dimension must be specified. If 'sparse', the vector dimension should not be specified. #[serde(rename = "vector_type", skip_serializing_if = "Option::is_none")] - pub vector_type: Option, + pub vector_type: Option, /// Identifies the name of the text field from your document model that is embedded. #[serde(rename = "field_map", skip_serializing_if = "Option::is_none")] pub field_map: Option, diff --git a/src/pinecone/control.rs b/src/pinecone/control.rs index fc49f9b..b1a98ba 100644 --- a/src/pinecone/control.rs +++ b/src/pinecone/control.rs @@ -1,15 +1,16 @@ use std::cmp::min; +use std::collections::HashMap; use std::time::Duration; use crate::openapi::apis::manage_indexes_api; -use crate::openapi::models::CreateIndexRequest; +use crate::openapi::models::{ByocSpec, CreateIndexForModelRequestEmbed, CreateIndexRequest}; use crate::pinecone::PineconeClient; use crate::utils::errors::PineconeError; use crate::models::{ - Cloud, CollectionList, CollectionModel, ConfigureIndexRequest, ConfigureIndexRequestSpec, - ConfigureIndexRequestSpecPod, CreateCollectionRequest, DeletionProtection, IndexList, - IndexModel, IndexSpec, Metric, PodSpec, PodSpecMetadataConfig, ServerlessSpec, WaitPolicy, + Cloud, ConfigureIndexRequest, ConfigureIndexRequestSpec, ConfigureIndexRequestSpecPod, + CreateIndexForModelOptions, DeletionProtection, IndexList, IndexModel, IndexSpec, Metric, + ServerlessSpec, VectorType, WaitPolicy, }; use crate::openapi::models; @@ -30,7 +31,7 @@ impl PineconeClient { /// /// ### Example /// ```no_run - /// use pinecone_sdk::models::{IndexModel, Metric, Cloud, WaitPolicy, DeletionProtection}; + /// use pinecone_sdk::models::{IndexModel, Metric, Cloud, WaitPolicy, DeletionProtection, VectorType}; /// use pinecone_sdk::utils::errors::PineconeError; /// /// # #[tokio::main] @@ -45,9 +46,9 @@ impl PineconeClient { /// Cloud::Aws, // Cloud provider /// "us-east-1", // Region /// DeletionProtection::Enabled, // Deletion protection - /// WaitPolicy::NoWait // Timeout + /// WaitPolicy::NoWait, // Timeout + /// VectorType::Dense, // Vector type /// None, - /// "dense".to_string(), // Vector type /// ).await; /// /// # Ok(()) @@ -63,13 +64,13 @@ impl PineconeClient { region: &str, deletion_protection: DeletionProtection, timeout: WaitPolicy, - tags: Option>, - vector_type: String, + vector_type: VectorType, + tags: Option>, ) -> Result { // create request specs let create_index_request_spec = IndexSpec { serverless: Some(Box::new(ServerlessSpec { - cloud, + cloud: cloud.into(), region: region.to_string(), })), pod: None, @@ -82,8 +83,8 @@ impl PineconeClient { deletion_protection: Some(deletion_protection), metric: Some(metric.into()), spec: Some(Box::new(create_index_request_spec)), - tags, vector_type: Some(vector_type), + tags, }; // make openAPI call @@ -98,20 +99,14 @@ impl PineconeClient { } } - /// Creates a pod index. + /// Creates a BYOC index. /// /// ### Arguments /// * `name: &str` - The name of the index /// * `dimension: i32` - The dimension of the index /// * `metric: Metric` - The metric to use for the index /// * `environment: &str` - The environment where the pod index will be deployed. Example: 'us-east1-gcp' - /// * `pod_type: &str` - This value combines pod type and pod size into a single string. This configuration is your main lever for vertical scaling. - /// * `pods: i32` - The number of pods to deploy. - /// * `replicas: i32` - The number of replicas to deploy for the pod index. - /// * `shards: i32` - The number of shards to use. Shards are used to expand the amount of vectors you can store beyond the capacity of a single pod. /// * `deletion_protection: DeletionProtection` - Deletion protection for the index. - /// * `metadata_indexed: Option<&[&str]>` - The metadata fields to index. - /// * `source_collection: Option<&str>` - The name of the collection to use as the source for the pod index. This configuration is only used when creating a pod index from an existing collection. /// * `timeout: WaitPolicy` - The wait policy for index creation. If the index becomes ready before the specified duration, the function will return early. If the index is not ready after the specified duration, the function will return an error. /// /// ### Return @@ -119,7 +114,7 @@ impl PineconeClient { /// /// ### Example /// ```no_run - /// use pinecone_sdk::models::{IndexModel, Metric, Cloud, WaitPolicy, DeletionProtection}; + /// use pinecone_sdk::models::{IndexModel, Metric, Cloud, WaitPolicy, DeletionProtection, VectorType}; /// use pinecone_sdk::utils::errors::PineconeError; /// use std::time::Duration; /// @@ -128,64 +123,41 @@ impl PineconeClient { /// let pinecone = pinecone_sdk::pinecone::default_client()?; /// /// // Create a pod index. - /// let response: Result = pinecone.create_pod_index( + /// let response: Result = pinecone.create_byoc_index( /// "index_name", // Name of the index /// 10, // Dimension of the index /// Metric::Cosine, // Distance metric - /// "us-east-1", // Environment - /// "p1.x1", // Pod type - /// 1, // Number of pods - /// 1, // Number of replicas - /// 1, // Number of shards + /// "aws-us-east-1-b921", // Environment /// DeletionProtection::Enabled, // Deletion protection - /// Some( // Metadata fields to index - /// &vec!["genre", - /// "title", - /// "imdb_rating"]), - /// Some("example-collection"), // Source collection /// WaitPolicy::WaitFor(Duration::from_secs(10)), // Timeout - /// None, - /// "dense".to_string(), // Vector type + /// VectorType::Dense, // Vector type + /// None, // tags /// ) /// .await; /// # Ok(()) /// # } /// ``` #[allow(clippy::too_many_arguments)] - pub async fn create_pod_index( + pub async fn create_byoc_index( &self, name: &str, dimension: i32, metric: Metric, environment: &str, - pod_type: &str, - pods: i32, - replicas: i32, - shards: i32, deletion_protection: DeletionProtection, - metadata_indexed: Option<&[&str]>, - source_collection: Option<&str>, timeout: WaitPolicy, - tags: Option>, - vector_type: String, + vector_type: VectorType, + tags: Option>, ) -> Result { // create request specs - let indexed = metadata_indexed.map(|i| i.iter().map(|s| s.to_string()).collect()); - - let pod_spec = PodSpec { + let spec = ByocSpec { environment: environment.to_string(), - replicas: Some(replicas), - shards: Some(shards), - pod_type: pod_type.to_string(), - pods: Some(pods), - metadata_config: Some(Box::new(PodSpecMetadataConfig { indexed })), - source_collection: source_collection.map(|s| s.to_string()), }; let spec = IndexSpec { serverless: None, - pod: Some(Box::new(pod_spec)), - byoc: None, + pod: None, + byoc: Some(Box::new(spec)), }; let create_index_request = CreateIndexRequest { @@ -194,8 +166,8 @@ impl PineconeClient { deletion_protection: Some(deletion_protection), metric: Some(metric.into()), spec: Some(Box::new(spec)), - tags, vector_type: Some(vector_type), + tags, }; // make openAPI call @@ -210,6 +182,40 @@ impl PineconeClient { } } + /// Creates an integrated index for a model. + #[allow(clippy::too_many_arguments)] + pub async fn create_index_for_model( + &self, + name: &str, + cloud: Cloud, + region: &str, + embed: CreateIndexForModelOptions, + deletion_protection: Option, + tags: Option>, + timeout: WaitPolicy, + ) -> Result { + let embed: CreateIndexForModelRequestEmbed = embed.into(); + + let request = models::CreateIndexForModelRequest { + name: name.to_string(), + cloud: cloud.into(), + region: region.to_string(), + embed: Box::new(embed), + deletion_protection, + tags, + }; + + let res = manage_indexes_api::create_index_for_model(&self.openapi_config, request) + .await + .map_err(PineconeError::from)?; + + // poll index status + match self.handle_poll_index(name, timeout).await { + Ok(_) => Ok(res.into()), + Err(e) => Err(e), + } + } + // Checks if the index is ready by polling the index status async fn handle_poll_index( &self, @@ -362,7 +368,7 @@ impl PineconeClient { deletion_protection: Option, replicas: Option, pod_type: Option<&str>, - tags: Option>, + tags: Option>, embed: Option>, ) -> Result { if replicas.is_none() && pod_type.is_none() && deletion_protection.is_none() { @@ -441,138 +447,6 @@ impl PineconeClient { Ok(()) } - - /// Creates a collection from an index. - /// - /// ### Arguments - /// * `name: &str` - Name of the collection to create. - /// * `source: &str` - Name of the index to be used as the source for the collection. - /// - /// ### Return - /// * `Result` - /// - /// ### Example - /// ```no_run - /// use pinecone_sdk::models::CollectionModel; - /// use pinecone_sdk::utils::errors::PineconeError; - /// - /// # #[tokio::main] - /// # async fn main() -> Result<(), PineconeError>{ - /// let pinecone = pinecone_sdk::pinecone::default_client()?; - /// - /// // Describe an index in the project. - /// let response: Result = pinecone.create_collection("collection-name", "index-name").await; - /// # Ok(()) - /// # } - /// ``` - pub async fn create_collection( - &self, - name: &str, - source: &str, - ) -> Result { - let create_collection_request = CreateCollectionRequest { - name: name.to_string(), - source: source.to_string(), - }; - - // make openAPI call - let res = - manage_indexes_api::create_collection(&self.openapi_config, create_collection_request) - .await - .map_err(PineconeError::from)?; - - Ok(res) - } - - /// Describe a collection. - /// - /// ### Arguments - /// * `name: &str` - The name of the collection to describe. - /// - /// ### Return - /// * `Result<(), PineconeError>` - /// - /// ### Example - /// ```no_run - /// use pinecone_sdk::models::CollectionModel; - /// use pinecone_sdk::utils::errors::PineconeError; - /// - /// # #[tokio::main] - /// # async fn main() -> Result<(), PineconeError>{ - /// let pinecone = pinecone_sdk::pinecone::default_client()?; - /// - /// // Describe a collection in the project. - /// let collection: CollectionModel = pinecone.describe_collection("collection-name").await?; - /// # Ok(()) - /// # } - /// ``` - pub async fn describe_collection(&self, name: &str) -> Result { - let res = manage_indexes_api::describe_collection(&self.openapi_config, name) - .await - .map_err(PineconeError::from)?; - - Ok(res) - } - - /// Lists all collections. - /// - /// This operation returns a list of all collections in a project. - /// - /// ### Return - /// * `Result` - /// - /// ### Example - /// ```no_run - /// use pinecone_sdk::models::CollectionList; - /// use pinecone_sdk::utils::errors::PineconeError; - /// - /// # #[tokio::main] - /// # async fn main() -> Result<(), PineconeError>{ - /// let pinecone = pinecone_sdk::pinecone::default_client()?; - /// - /// // List all collections in the project. - /// let response: Result = pinecone.list_collections().await; - /// # Ok(()) - /// # } - /// ``` - pub async fn list_collections(&self) -> Result { - // make openAPI call - let res = manage_indexes_api::list_collections(&self.openapi_config) - .await - .map_err(PineconeError::from)?; - - Ok(res) - } - - /// Deletes a collection. - /// - /// ### Arguments - /// * `name: &str` - The name of the collection to be deleted. - /// - /// ### Return - /// * `Result<(), PineconeError>` - /// - /// ### Example - /// ```no_run - /// use pinecone_sdk::utils::errors::PineconeError; - /// - /// # #[tokio::main] - /// # async fn main() -> Result<(), PineconeError>{ - /// let pinecone = pinecone_sdk::pinecone::default_client()?; - /// - /// // Delete a collection in the project. - /// let response: Result<(), PineconeError> = pinecone.delete_collection("collection-name").await; - /// # Ok(()) - /// # } - /// ``` - pub async fn delete_collection(&self, name: &str) -> Result<(), PineconeError> { - // make openAPI call - manage_indexes_api::delete_collection(&self.openapi_config, name) - .await - .map_err(PineconeError::from)?; - - Ok(()) - } } #[cfg(test)] @@ -580,7 +454,7 @@ mod tests { use super::*; use crate::openapi::{ self, - models::{self, collection_model::Status}, + models::{self}, }; use crate::pinecone::PineconeClientConfig; use httpmock::prelude::*; @@ -632,8 +506,8 @@ mod tests { "us-east-1", DeletionProtection::Enabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create serverless index"); @@ -696,8 +570,8 @@ mod tests { "us-east-1", DeletionProtection::Enabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create serverless index"); @@ -750,8 +624,8 @@ mod tests { "abc", DeletionProtection::Enabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect_err("Expected error when creating serverless index"); @@ -800,8 +674,8 @@ mod tests { "us-west-1", DeletionProtection::Enabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect_err("Expected error when creating serverless index"); @@ -824,14 +698,14 @@ mod tests { then.status(422) .header("content-type", "application/json") .body( - r#"{ + r#"{ "error": { "code": "INVALID_ARGUMENT", "message": "Failed to deserialize the JSON body into the target type: missing field `metric` at line 1 column 16" }, "status": 422 }"#, - ); + ); }); let config = PineconeClientConfig { @@ -850,8 +724,8 @@ mod tests { "us-west-1", DeletionProtection::Enabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect_err("Expected error when creating serverless index"); @@ -890,8 +764,8 @@ mod tests { "us-east-1", DeletionProtection::Enabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect_err("Expected create_index to return an error"); @@ -967,6 +841,9 @@ mod tests { pod: None, byoc: None, }, + vector_type: VectorType::Dense, + tags: None, + embed: None, }; assert_eq!(index, expected); @@ -1107,6 +984,9 @@ mod tests { deletion_protection: None, spec: models::IndexModelSpec::default(), status: models::IndexModelStatus::default(), + embed: None, + tags: None, + vector_type: VectorType::Dense, }, IndexModel { name: "index2".to_string(), @@ -1116,6 +996,9 @@ mod tests { deletion_protection: None, spec: models::IndexModelSpec::default(), status: models::IndexModelStatus::default(), + embed: None, + tags: None, + vector_type: VectorType::Dense, }, ]), }; @@ -1156,43 +1039,32 @@ mod tests { } #[tokio::test] - async fn test_create_pod_index() -> Result<(), PineconeError> { + async fn test_handle_polling_index_ok() -> Result<(), PineconeError> { let server = MockServer::start(); let mock = server.mock(|when, then| { - when.method(POST).path("/indexes"); - then.status(201) + when.method(GET).path("/indexes/index-name"); + then.status(200) .header("content-type", "application/json") .body( r#" { - "name": "index-name", "dimension": 1536, - "metric": "euclidean", "host": "mock-host", + "metric": "cosine", + "name": "index-name", "spec": { - "pod": { - "environment": "us-east-1-aws", - "replicas": 1, - "shards": 1, - "pod_type": "p1.x1", - "pods": 1, - "metadata_config": { - "indexed": [ - "genre", - "title", - "imdb_rating" - ] - } + "serverless": { + "cloud": "aws", + "region": "us-east-1" } }, "status": { "ready": true, - "state": "ScalingUpPodSize" + "state": "Ready" }, "vector_type": "dense" - } - "#, + }"#, ); }); @@ -1203,57 +1075,73 @@ mod tests { }; let pinecone = config.client().expect("Failed to create Pinecone instance"); - let create_index_response = pinecone - .create_pod_index( - "index-name", - 1536, - Metric::Euclidean, - "us-east-1-aws", - "p1.x1", - 1, - 1, - 1, - DeletionProtection::Enabled, - Some(&["genre", "title", "imdb_rating"]), - Some("example-collection"), - WaitPolicy::NoWait, - None, - "dense".to_string(), - ) - .await - .expect("Failed to create pod index"); + let res = pinecone + .handle_poll_index("index-name", WaitPolicy::WaitFor(Duration::from_secs(1))) + .await; - assert_eq!(create_index_response.name, "index-name"); - assert_eq!(create_index_response.dimension, Some(1536)); - assert_eq!(create_index_response.metric, Metric::Euclidean); + assert!(res.is_ok()); + mock.assert(); - let pod_spec = create_index_response.spec.pod.as_ref().unwrap(); - assert_eq!(pod_spec.environment, "us-east-1-aws"); - assert_eq!(pod_spec.pod_type, "p1.x1"); - assert_eq!( - pod_spec.metadata_config.as_ref().unwrap().indexed, - Some(vec![ - "genre".to_string(), - "title".to_string(), - "imdb_rating".to_string() - ]) - ); - assert_eq!(pod_spec.pods, Some(1)); - assert_eq!(pod_spec.replicas, Some(1)); - assert_eq!(pod_spec.shards, Some(1)); + Ok(()) + } - mock.assert(); + #[tokio::test] + async fn test_handle_polling_index_err() -> Result<(), PineconeError> { + let server = MockServer::start(); + + let mock = server.mock(|when, then| { + when.method(GET).path("/indexes/index-name"); + then.status(200) + .header("content-type", "application/json") + .body( + r#" + { + "dimension": 1536, + "host": "mock-host", + "metric": "cosine", + "name": "index-name", + "spec": { + "serverless": { + "cloud": "aws", + "region": "us-east-1" + } + }, + "status": { + "ready": false, + "state": "Initializing" + } + }"#, + ); + }); + + let config = PineconeClientConfig { + api_key: Some("api_key".to_string()), + control_plane_host: Some(server.base_url()), + ..Default::default() + }; + let pinecone = config.client().expect("Failed to create Pinecone instance"); + + let start_time = std::time::Instant::now(); + let err = pinecone + .handle_poll_index("index-name", WaitPolicy::WaitFor(Duration::from_secs(7))) + .await + .expect_err("Expected to fail polling index"); + + assert!(start_time.elapsed().as_secs() >= 7 && start_time.elapsed().as_secs() < 8); + assert!(matches!(err, PineconeError::TimeoutError { .. })); + + mock.assert_hits(3); Ok(()) } #[tokio::test] - async fn test_create_pod_index_with_defaults() -> Result<(), PineconeError> { + async fn test_configure_index() -> Result<(), PineconeError> { let server = MockServer::start(); let mock = server.mock(|when, then| { - when.method(POST).path("/indexes"); - then.status(201) + when.path("/indexes/index-name"); + then.status(202) .header("content-type", "application/json") .body( r#" @@ -1265,11 +1153,17 @@ mod tests { "spec": { "pod": { "environment": "us-east-1-aws", + "replicas": 6, + "shards": 1, "pod_type": "p1.x1", "pods": 1, - "metadata_config": {}, - "replicas": 1, - "shards": 1 + "metadata_config": { + "indexed": [ + "genre", + "title", + "imdb_rating" + ] + } } }, "status": { @@ -1277,8 +1171,7 @@ mod tests { "state": "ScalingUpPodSize" }, "vector_type": "dense" - } - "#, + }"#, ); }); @@ -1289,37 +1182,23 @@ mod tests { }; let pinecone = config.client().expect("Failed to create Pinecone instance"); - let create_index_response = pinecone - .create_pod_index( + let configure_index_response = pinecone + .configure_index( "index-name", - 1536, - Default::default(), - "us-east-1-aws", - "p1.x1", - 1, - 1, - 1, - DeletionProtection::Enabled, - None, + Some(DeletionProtection::Disabled), + Some(6), + Some("p1.x1"), None, - WaitPolicy::NoWait, None, - "dense".to_string(), ) .await - .expect("Failed to create pod index"); + .expect("Failed to configure index"); - assert_eq!(create_index_response.name, "index-name"); - assert_eq!(create_index_response.dimension, Some(1536)); - assert_eq!(create_index_response.metric, Metric::Cosine); + assert_eq!(configure_index_response.name, "index-name"); - let pod_spec = create_index_response.spec.pod.as_ref().unwrap(); - assert_eq!(pod_spec.environment, "us-east-1-aws"); - assert_eq!(pod_spec.pod_type, "p1.x1"); - assert_eq!(pod_spec.metadata_config.as_ref().unwrap().indexed, None); - assert_eq!(pod_spec.pods, Some(1)); - assert_eq!(pod_spec.replicas, Some(1)); - assert_eq!(pod_spec.shards, Some(1)); + let spec = configure_index_response.spec.pod.unwrap(); + assert_eq!(spec.replicas, Some(6)); + assert_eq!(spec.pod_type.as_str(), "p1.x1"); mock.assert(); @@ -1327,760 +1206,45 @@ mod tests { } #[tokio::test] - async fn test_create_pod_index_quota_exceeded() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(POST).path("/indexes"); - then.status(403) - .header("content-type", "application/json") - .body( - r#" - { - "error": { - "code": "FORBIDDEN", - "message": "Increase yoru quota or upgrade to create more indexes." - }, - "status": 403 - } - "#, - ); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let create_index_response = pinecone - .create_pod_index( - "index-name", - 1536, - Metric::Euclidean, - "test-environment", - "p1.x1", - 1, - 1, - 1, - DeletionProtection::Enabled, - None, - Some("example-collection"), - WaitPolicy::NoWait, - None, - "dense".to_string(), - ) - .await - .expect_err("Expected create_pod_index to return an error"); - - assert!(matches!( - create_index_response, - PineconeError::PodQuotaExceededError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_create_pod_index_invalid_environment() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(POST).path("/indexes"); - then.status(400) - .header("content-type", "application/json") - .body( - r#" - { - "error": "Invalid environment" - } - "#, - ); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let create_index_response = pinecone - .create_pod_index( - "index-name", - 1536, - Metric::Euclidean, - "invalid-environment", - "p1.x1", - 1, - 1, - 1, - DeletionProtection::Enabled, - Some(&["genre", "title", "imdb_rating"]), - Some("example-collection"), - WaitPolicy::NoWait, - None, - "dense".to_string(), - ) - .await - .expect_err("Expected create_pod_index to return an error"); - - assert!(matches!( - create_index_response, - PineconeError::BadRequestError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_create_pod_index_invalid_pod_type() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(POST).path("/indexes"); - then.status(400) - .header("content-type", "application/json") - .body( - r#" - { - "error": "Invalid pod type" - } - "#, - ); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let create_index_response = pinecone - .create_pod_index( - "index-name", - 1536, - Metric::Euclidean, - "us-east-1-aws", - "invalid-pod-type", - 1, - 1, - 1, - DeletionProtection::Enabled, - Some(&["genre", "title", "imdb_rating"]), - Some("example-collection"), - WaitPolicy::NoWait, - None, - "dense".to_string(), - ) - .await - .expect_err("Expected create_pod_index to return an error"); - - assert!(matches!( - create_index_response, - PineconeError::BadRequestError { .. } - )); - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_handle_polling_index_ok() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(GET).path("/indexes/index-name"); - then.status(200) - .header("content-type", "application/json") - .body( - r#" - { - "dimension": 1536, - "host": "mock-host", - "metric": "cosine", - "name": "index-name", - "spec": { - "serverless": { - "cloud": "aws", - "region": "us-east-1" - } - }, - "status": { - "ready": true, - "state": "Ready" - }, - "vector_type": "dense" - }"#, - ); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let res = pinecone - .handle_poll_index("index-name", WaitPolicy::WaitFor(Duration::from_secs(1))) - .await; - - assert!(res.is_ok()); - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_handle_polling_index_err() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(GET).path("/indexes/index-name"); - then.status(200) - .header("content-type", "application/json") - .body( - r#" - { - "dimension": 1536, - "host": "mock-host", - "metric": "cosine", - "name": "index-name", - "spec": { - "serverless": { - "cloud": "aws", - "region": "us-east-1" - } - }, - "status": { - "ready": false, - "state": "Initializing" - } - }"#, - ); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let start_time = std::time::Instant::now(); - let err = pinecone - .handle_poll_index("index-name", WaitPolicy::WaitFor(Duration::from_secs(7))) - .await - .expect_err("Expected to fail polling index"); - - assert!(start_time.elapsed().as_secs() >= 7 && start_time.elapsed().as_secs() < 8); - assert!(matches!(err, PineconeError::TimeoutError { .. })); - - mock.assert_hits(3); - - Ok(()) - } - - #[tokio::test] - async fn test_configure_index() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.path("/indexes/index-name"); - then.status(202) - .header("content-type", "application/json") - .body( - r#" - { - "name": "index-name", - "dimension": 1536, - "metric": "cosine", - "host": "mock-host", - "spec": { - "pod": { - "environment": "us-east-1-aws", - "replicas": 6, - "shards": 1, - "pod_type": "p1.x1", - "pods": 1, - "metadata_config": { - "indexed": [ - "genre", - "title", - "imdb_rating" - ] - } - } - }, - "status": { - "ready": true, - "state": "ScalingUpPodSize" - }, - "vector_type": "dense" - }"#, - ); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let configure_index_response = pinecone - .configure_index( - "index-name", - Some(DeletionProtection::Disabled), - Some(6), - Some("p1.x1"), - None, - None, - ) - .await - .expect("Failed to configure index"); - - assert_eq!(configure_index_response.name, "index-name"); - - let spec = configure_index_response.spec.pod.unwrap(); - assert_eq!(spec.replicas, Some(6)); - assert_eq!(spec.pod_type.as_str(), "p1.x1"); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_configure_deletion_protection() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.path("/indexes/index-name"); - then.status(202) - .header("content-type", "application/json") - .body( - r#"{ - "name": "index-name", - "dimension": 1536, - "metric": "cosine", - "host": "mock-host", - "deletion_protection": "disabled", - "spec": { - "pod": { - "environment": "us-east-1-aws", - "metadata_config": { - "indexed": [ - "genre", - "title", - "imdb_rating" - ] - }, - "pod_type": "p1.x1", - "pods": 1, - "replicas": 1, - "shards": 1 - } - }, - "status": { - "ready": true, - "state": "ScalingUpPodSize" - }, - "vector_type": "dense" - }"#, - ); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let configure_index_response = pinecone - .configure_index( - "index-name", - Some(DeletionProtection::Disabled), - None, - None, - None, - None, - ) - .await - .expect("Failed to configure index"); - - assert_eq!( - configure_index_response.deletion_protection, - Some(DeletionProtection::Disabled) - ); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_configure_index_no_config() -> Result<(), PineconeError> { - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let configure_index_response = pinecone - .configure_index("index-name", None, None, None, None, None) - .await; - - assert!(matches!( - configure_index_response, - Err(PineconeError::InvalidConfigurationError { .. }) - )); - - Ok(()) - } - - #[tokio::test] - async fn test_configure_index_quota_exceeded() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.path("/indexes/index-name"); - then.status(403) - .header("content-type", "application/json") - .body( - r#" - { - "error": { - "code": "FORBIDDEN", - "message": "Increase your quota or upgrade to create more indexes." - }, - "status": 403 - } - "#, - ); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let configure_index_response = pinecone - .configure_index( - "index-name", - Some(DeletionProtection::Enabled), - Some(6), - Some("p1.x1"), - None, - None, - ) - .await - .expect_err("Expected to fail to configure index"); - - assert!(matches!( - configure_index_response, - PineconeError::PodQuotaExceededError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_configure_index_not_found() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.path("/indexes/index-name"); - then.status(404) - .header("content-type", "application/json") - .body( - r#"{ - "error": { - "code": "NOT_FOUND", - "message": "Index index-name not found." - }, - "status": 404 - }"#, - ); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let configure_index_response = pinecone - .configure_index( - "index-name", - Some(DeletionProtection::Disabled), - Some(6), - Some("p1.x1"), - None, - None, - ) - .await - .expect_err("Expected to fail to configure index"); - - assert!(matches!( - configure_index_response, - PineconeError::IndexNotFoundError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_configure_index_unprocessable_entity() -> Result<(), PineconeError> { + async fn test_configure_deletion_protection() -> Result<(), PineconeError> { let server = MockServer::start(); let mock = server.mock(|when, then| { when.path("/indexes/index-name"); - then.status(422) + then.status(202) .header("content-type", "application/json") .body( r#"{ - "error": { - "code": "INVALID_ARGUMENT", - "message": "Failed to deserialize the JSON body into the target type - }, - "status": 422 - }"#, - ); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let configure_index_response = pinecone - .configure_index( - "index-name", - Some(DeletionProtection::Enabled), - Some(6), - Some("p1.x1"), - None, - None, - ) - .await - .expect_err("Expected to fail to configure index"); - - assert!(matches!( - configure_index_response, - PineconeError::UnprocessableEntityError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_configure_index_internal_error() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.path("/indexes/index-name"); - then.status(500); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let configure_index_response = pinecone - .configure_index( - "index-name", - Some(DeletionProtection::Enabled), - Some(6), - Some("p1.x1"), - None, - None, - ) - .await - .expect_err("Expected to fail to configure index"); - - assert!(matches!( - configure_index_response, - PineconeError::InternalServerError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_delete_index() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(DELETE).path("/indexes/index-name"); - then.status(202); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - pinecone - .delete_index("index-name") - .await - .expect("Failed to delete index"); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_delete_index_invalid_name() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(DELETE).path("/indexes/invalid-index"); - then.status(404) - .header("content-type", "application/json") - .body( - r#" - { - "error": "Index not found" - } - "#, - ); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let delete_index_response = pinecone - .delete_index("invalid-index") - .await - .expect_err("Expected delete_index to return an error"); - - assert!(matches!( - delete_index_response, - PineconeError::IndexNotFoundError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_delete_index_pending_collection() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(DELETE).path("/indexes/index-name"); - then.status(412); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let delete_index_response = pinecone - .delete_index("index-name") - .await - .expect_err("Expected delete_index to return an error"); - - assert!(matches!( - delete_index_response, - PineconeError::PendingCollectionError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_delete_index_server_error() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(DELETE).path("/indexes/index-name"); - then.status(500); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let delete_index_response = pinecone - .delete_index("index-name") - .await - .expect_err("Expected delete_index to return an error"); - - assert!(matches!( - delete_index_response, - PineconeError::InternalServerError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_create_collection() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(POST).path("/collections"); - then.status(201) - .header("content-type", "application/json") - .body( - r#" - { - "name": "example-collection", - "size": 10000000, - "status": "Initializing", + "name": "index-name", "dimension": 1536, - "vector_count": 120000, - "environment": "us-east1-gcp" - } - "#, + "metric": "cosine", + "host": "mock-host", + "deletion_protection": "disabled", + "spec": { + "pod": { + "environment": "us-east-1-aws", + "metadata_config": { + "indexed": [ + "genre", + "title", + "imdb_rating" + ] + }, + "pod_type": "p1.x1", + "pods": 1, + "replicas": 1, + "shards": 1 + } + }, + "status": { + "ready": true, + "state": "ScalingUpPodSize" + }, + "vector_type": "dense" + }"#, ); }); - // Construct Pinecone instance with the mock server URL let config = PineconeClientConfig { api_key: Some("api_key".to_string()), control_plane_host: Some(server.base_url()), @@ -2088,21 +1252,22 @@ mod tests { }; let pinecone = config.client().expect("Failed to create Pinecone instance"); - // Call create_collection and verify the result - let collection = pinecone - .create_collection("collection1", "index1") + let configure_index_response = pinecone + .configure_index( + "index-name", + Some(DeletionProtection::Disabled), + None, + None, + None, + None, + ) .await - .expect("Failed to create collection"); + .expect("Failed to configure index"); - let expected = CollectionModel { - name: "example-collection".to_string(), - size: Some(10000000), - status: Status::Initializing, - dimension: Some(1536), - vector_count: Some(120000), - environment: "us-east1-gcp".to_string(), - }; - assert_eq!(collection, expected); + assert_eq!( + configure_index_response.deletion_protection, + Some(DeletionProtection::Disabled) + ); mock.assert(); @@ -2110,60 +1275,41 @@ mod tests { } #[tokio::test] - async fn test_create_collection_quota_exceeded() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(POST).path("/collections"); - then.status(403) - .header("content-type", "application/json") - .body( - r#" - { - "error": { - "code": "FORBIDDEN", - "message": "Collection exceeds quota. Maximum allowed on your account is 1. Currently have 1." - }, - "status": 403 - } - "#, - ); - }); - + async fn test_configure_index_no_config() -> Result<(), PineconeError> { let config = PineconeClientConfig { api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), ..Default::default() }; let pinecone = config.client().expect("Failed to create Pinecone instance"); - let create_collection_response = pinecone - .create_collection("invalid_collection", "valid-index") - .await - .expect_err("Expected create_collection to return an error"); + let configure_index_response = pinecone + .configure_index("index-name", None, None, None, None, None) + .await; assert!(matches!( - create_collection_response, - PineconeError::CollectionsQuotaExceededError { .. } + configure_index_response, + Err(PineconeError::InvalidConfigurationError { .. }) )); - mock.assert(); - Ok(()) } #[tokio::test] - async fn test_create_collection_invalid_name() -> Result<(), PineconeError> { + async fn test_configure_index_quota_exceeded() -> Result<(), PineconeError> { let server = MockServer::start(); let mock = server.mock(|when, then| { - when.method(POST).path("/collections"); - then.status(409) + when.path("/indexes/index-name"); + then.status(403) .header("content-type", "application/json") .body( r#" { - "error": "Index not found" + "error": { + "code": "FORBIDDEN", + "message": "Increase your quota or upgrade to create more indexes." + }, + "status": 403 } "#, ); @@ -2176,45 +1322,21 @@ mod tests { }; let pinecone = config.client().expect("Failed to create Pinecone instance"); - let create_collection_response = pinecone - .create_collection("invalid_collection", "valid-index") - .await - .expect_err("Expected create_collection to return an error"); - - assert!(matches!( - create_collection_response, - PineconeError::ResourceAlreadyExistsError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_create_collection_server_error() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(POST).path("/collections"); - then.status(500); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - let create_collection_response = pinecone - .create_collection("collection-name", "index1") + let configure_index_response = pinecone + .configure_index( + "index-name", + Some(DeletionProtection::Enabled), + Some(6), + Some("p1.x1"), + None, + None, + ) .await - .expect_err("Expected create_collection to return an error"); + .expect_err("Expected to fail to configure index"); assert!(matches!( - create_collection_response, - PineconeError::InternalServerError { .. } + configure_index_response, + PineconeError::PodQuotaExceededError { .. } )); mock.assert(); @@ -2223,26 +1345,24 @@ mod tests { } #[tokio::test] - async fn test_describe_collection() -> Result<(), PineconeError> { + async fn test_configure_index_not_found() -> Result<(), PineconeError> { let server = MockServer::start(); let mock = server.mock(|when, then| { - when.method(GET).path("/collections/collection-name"); - then.status(200) + when.path("/indexes/index-name"); + then.status(404) .header("content-type", "application/json") .body( r#"{ - "dimension": 3, - "environment": "us-east1-gcp", - "name": "tiny-collection", - "size": 3126700, - "status": "Ready", - "vector_count": 99 - }"#, + "error": { + "code": "NOT_FOUND", + "message": "Index index-name not found." + }, + "status": 404 + }"#, ); }); - // Construct Pinecone instance with the mock server URL let config = PineconeClientConfig { api_key: Some("api_key".to_string()), control_plane_host: Some(server.base_url()), @@ -2250,38 +1370,43 @@ mod tests { }; let pinecone = config.client().expect("Failed to create Pinecone instance"); - // Call describe_collection and verify the result - let collection = pinecone - .describe_collection("collection-name") + let configure_index_response = pinecone + .configure_index( + "index-name", + Some(DeletionProtection::Disabled), + Some(6), + Some("p1.x1"), + None, + None, + ) .await - .expect("Failed to describe collection"); - - let expected = CollectionModel { - name: "tiny-collection".to_string(), - size: Some(3126700), - status: Status::Ready, - dimension: Some(3), - vector_count: Some(99), - environment: "us-east1-gcp".to_string(), - }; + .expect_err("Expected to fail to configure index"); + + assert!(matches!( + configure_index_response, + PineconeError::IndexNotFoundError { .. } + )); - assert_eq!(collection, expected); mock.assert(); Ok(()) } #[tokio::test] - async fn test_describe_collection_invalid_name() -> Result<(), PineconeError> { + async fn test_configure_index_unprocessable_entity() -> Result<(), PineconeError> { let server = MockServer::start(); let mock = server.mock(|when, then| { - when.method(GET).path("/collections/invalid-collection"); - then.status(404) + when.path("/indexes/index-name"); + then.status(422) .header("content-type", "application/json") .body( r#"{ - "error": "Collection invalid-collection not found" + "error": { + "code": "INVALID_ARGUMENT", + "message": "Failed to deserialize the JSON body into the target type + }, + "status": 422 }"#, ); }); @@ -2293,26 +1418,34 @@ mod tests { }; let pinecone = config.client().expect("Failed to create Pinecone instance"); - let response = pinecone - .describe_collection("invalid-collection") + let configure_index_response = pinecone + .configure_index( + "index-name", + Some(DeletionProtection::Enabled), + Some(6), + Some("p1.x1"), + None, + None, + ) .await - .expect_err("Expected describe_collection to return an error"); + .expect_err("Expected to fail to configure index"); assert!(matches!( - response, - PineconeError::CollectionNotFoundError { .. } + configure_index_response, + PineconeError::UnprocessableEntityError { .. } )); + mock.assert(); Ok(()) } #[tokio::test] - async fn test_describe_collection_server_error() -> Result<(), PineconeError> { + async fn test_configure_index_internal_error() -> Result<(), PineconeError> { let server = MockServer::start(); let mock = server.mock(|when, then| { - when.method(GET).path("/collections/collection-name"); + when.path("/indexes/index-name"); then.status(500); }); @@ -2323,62 +1456,37 @@ mod tests { }; let pinecone = config.client().expect("Failed to create Pinecone instance"); - let response = pinecone - .describe_collection("collection-name") + let configure_index_response = pinecone + .configure_index( + "index-name", + Some(DeletionProtection::Enabled), + Some(6), + Some("p1.x1"), + None, + None, + ) .await - .expect_err("Expected describe_collection to return an error"); + .expect_err("Expected to fail to configure index"); assert!(matches!( - response, + configure_index_response, PineconeError::InternalServerError { .. } )); + mock.assert(); Ok(()) } #[tokio::test] - async fn test_list_collections() -> Result<(), PineconeError> { + async fn test_delete_index() -> Result<(), PineconeError> { let server = MockServer::start(); let mock = server.mock(|when, then| { - when.method(GET).path("/collections"); - then.status(200) - .header("content-type", "application/json") - .body( - r#" - { - "collections": [ - { - "name": "small-collection", - "size": 3126700, - "status": "Ready", - "dimension": 3, - "vector_count": 99, - "environment": "us-east1-gcp" - }, - { - "name": "small-collection-new", - "size": 3126700, - "status": "Initializing", - "dimension": 3, - "vector_count": 99, - "environment": "us-east1-gcp" - }, - { - "name": "big-collection", - "size": 160087040000000, - "status": "Ready", - "dimension": 1536, - "vector_count": 10000000, - "environment": "us-east1-gcp" - } - ] - }"#, - ); + when.method(DELETE).path("/indexes/index-name"); + then.status(202); }); - // Construct Pinecone instance with the mock server URL let config = PineconeClientConfig { api_key: Some("api_key".to_string()), control_plane_host: Some(server.base_url()), @@ -2386,42 +1494,10 @@ mod tests { }; let pinecone = config.client().expect("Failed to create Pinecone instance"); - // Call list_collections and verify the result - let collection_list = pinecone - .list_collections() + pinecone + .delete_index("index-name") .await - .expect("Failed to list collections"); - - let expected = CollectionList { - // name: String, dimension: i32, metric: Metric, host: String, spec: models::IndexModelSpec, status: models::IndexModelStatus) - collections: Some(vec![ - CollectionModel { - name: "small-collection".to_string(), - size: Some(3126700), - status: Status::Ready, - dimension: Some(3), - vector_count: Some(99), - environment: "us-east1-gcp".to_string(), - }, - CollectionModel { - name: "small-collection-new".to_string(), - size: Some(3126700), - status: Status::Initializing, - dimension: Some(3), - vector_count: Some(99), - environment: "us-east1-gcp".to_string(), - }, - CollectionModel { - name: "big-collection".to_string(), - size: Some(160087040000000), - status: Status::Ready, - dimension: Some(1536), - vector_count: Some(10000000), - environment: "us-east1-gcp".to_string(), - }, - ]), - }; - assert_eq!(collection_list, expected); + .expect("Failed to delete index"); mock.assert(); @@ -2429,12 +1505,20 @@ mod tests { } #[tokio::test] - async fn test_list_collections_error() -> Result<(), PineconeError> { + async fn test_delete_index_invalid_name() -> Result<(), PineconeError> { let server = MockServer::start(); let mock = server.mock(|when, then| { - when.method(GET).path("/collections"); - then.status(500); + when.method(DELETE).path("/indexes/invalid-index"); + then.status(404) + .header("content-type", "application/json") + .body( + r#" + { + "error": "Index not found" + } + "#, + ); }); let config = PineconeClientConfig { @@ -2444,41 +1528,15 @@ mod tests { }; let pinecone = config.client().expect("Failed to create Pinecone instance"); - // Call list_collections and verify the result - let list_collections_response = pinecone - .list_collections() + let delete_index_response = pinecone + .delete_index("invalid-index") .await - .expect_err("Expected to fail to list collections"); + .expect_err("Expected delete_index to return an error"); assert!(matches!( - list_collections_response, - PineconeError::InternalServerError { .. } + delete_index_response, + PineconeError::IndexNotFoundError { .. } )); - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_delete_collection() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(DELETE).path("/collections/collection-name"); - then.status(202); - }); - - let config = PineconeClientConfig { - api_key: Some("api_key".to_string()), - control_plane_host: Some(server.base_url()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - pinecone - .delete_collection("collection-name") - .await - .expect("Failed to delete collection"); mock.assert(); @@ -2486,20 +1544,12 @@ mod tests { } #[tokio::test] - async fn test_delete_collection_not_found() -> Result<(), PineconeError> { + async fn test_delete_index_pending_collection() -> Result<(), PineconeError> { let server = MockServer::start(); let mock = server.mock(|when, then| { - when.method(DELETE).path("/collections/collection-name"); - then.status(404) - .header("content-type", "application/json") - .body( - r#" - { - "error": "Collection not found" - } - "#, - ); + when.method(DELETE).path("/indexes/index-name"); + then.status(412); }); let config = PineconeClientConfig { @@ -2509,14 +1559,14 @@ mod tests { }; let pinecone = config.client().expect("Failed to create Pinecone instance"); - let delete_collection_response = pinecone - .delete_collection("collection-name") + let delete_index_response = pinecone + .delete_index("index-name") .await - .expect_err("Expected delete_collection to return an error"); + .expect_err("Expected delete_index to return an error"); assert!(matches!( - delete_collection_response, - PineconeError::CollectionNotFoundError { .. } + delete_index_response, + PineconeError::PendingCollectionError { .. } )); mock.assert(); @@ -2525,11 +1575,11 @@ mod tests { } #[tokio::test] - async fn test_delete_collection_internal_error() -> Result<(), PineconeError> { + async fn test_delete_index_server_error() -> Result<(), PineconeError> { let server = MockServer::start(); let mock = server.mock(|when, then| { - when.method(DELETE).path("/collections/collection-name"); + when.method(DELETE).path("/indexes/index-name"); then.status(500); }); @@ -2540,13 +1590,13 @@ mod tests { }; let pinecone = config.client().expect("Failed to create Pinecone instance"); - let delete_collection_response = pinecone - .delete_collection("collection-name") + let delete_index_response = pinecone + .delete_index("index-name") .await - .expect_err("Expected delete_collection to return an error"); + .expect_err("Expected delete_index to return an error"); assert!(matches!( - delete_collection_response, + delete_index_response, PineconeError::InternalServerError { .. } )); diff --git a/src/pinecone/data.rs b/src/pinecone/data.rs index ac283d3..728dc2d 100644 --- a/src/pinecone/data.rs +++ b/src/pinecone/data.rs @@ -1,7 +1,8 @@ use crate::pinecone::PineconeClient; use crate::protos::vector_service_client::VectorServiceClient; -use crate::utils::errors::PineconeError; +use crate::utils::errors::{handle_response_error, PineconeError}; use once_cell::sync::Lazy; +use std::sync::Arc; use tonic::metadata::{Ascii, MetadataValue as TonicMetadataVal}; use tonic::service::interceptor::InterceptedService; use tonic::service::Interceptor; @@ -12,6 +13,13 @@ use crate::models::{ DescribeIndexStatsResponse, FetchResponse, ListResponse, Metadata, Namespace, QueryResponse, SparseValues, UpdateResponse, UpsertResponse, Vector, }; +use crate::openapi::apis::configuration::Configuration; +use crate::openapi::apis::vector_operations_api::UpsertRecordsNamespaceError; +use crate::openapi::apis::{vector_operations_api, ResponseContent}; +use crate::openapi::models::{ + SearchRecordsRequest, SearchRecordsRequestQuery, SearchRecordsRequestRerank, + SearchRecordsResponse, +}; use crate::protos; #[derive(Debug, Clone)] @@ -33,11 +41,9 @@ impl Interceptor for ApiKeyInterceptor { /// A client for interacting with a Pinecone index. #[derive(Debug)] -#[allow(dead_code)] pub struct Index { - /// The name of the index. - host: String, connection: VectorServiceClient>, + client: Arc, } impl Index { @@ -94,6 +100,107 @@ impl Index { Ok(response) } + /// TODO + pub async fn upsert_records( + &mut self, + namespace: &str, + records: &[serde_json::Value], + ) -> Result<(), PineconeError> { + let configuration = &self.client.openapi_config; + let client = &self.client.openapi_config.client; + + let uri_str = format!( + "{}/records/namespaces/{namespace}/upsert", + configuration.base_path, + namespace = crate::openapi::apis::urlencode(namespace) + ); + let mut req_builder = client.request(reqwest::Method::POST, uri_str.as_str()); + + if let Some(ref user_agent) = configuration.user_agent { + req_builder = req_builder.header(reqwest::header::USER_AGENT, user_agent.clone()); + } + if let Some(ref apikey) = configuration.api_key { + let key = apikey.key.clone(); + let value = match apikey.prefix { + Some(ref prefix) => format!("{} {}", prefix, key), + None => key, + }; + req_builder = req_builder.header("Api-Key", value); + }; + req_builder = req_builder.header(reqwest::header::CONTENT_TYPE, "application/x-ndjson"); + + let ndjson = records + .iter() + .map(|r| r.to_string()) + .collect::>() + .join("\n"); + + req_builder = req_builder.body(ndjson); + + let req = req_builder + .build() + .map_err(|e| PineconeError::ReqwestError { + source: anyhow::Error::from(e), + })?; + + let resp = client + .execute(req) + .await + .map_err(|e| PineconeError::ReqwestError { + source: anyhow::Error::from(e), + })?; + + let status = resp.status(); + let content = resp.text().await.map_err(|e| PineconeError::ReqwestError { + source: anyhow::Error::from(e), + })?; + + if !status.is_client_error() && !status.is_server_error() { + Ok(()) + } else { + let entity: Option = serde_json::from_str(&content).ok(); + Err(handle_response_error( + ResponseContent { + status, + content, + entity, + } + .into(), + )) + } + } + + /// TODO + pub async fn search_records_by_text( + &mut self, + namespace: &str, + query_text: &str, + top_k: u32, + fields: Option>, + rerank: Option, + ) -> Result { + let response = vector_operations_api::search_records_namespace( + &self.client.openapi_config, + namespace, + SearchRecordsRequest { + query: Box::new(SearchRecordsRequestQuery { + inputs: Some(serde_json::json!({ + "text": query_text, + })), + top_k: top_k as i32, + filter: None, + vector: None, + id: None, + }), + fields, + rerank: rerank.map(Box::new), + }, + ) + .await?; + + Ok(response) + } + /// The list operation lists the IDs of vectors in a single namespace of a serverless index. An optional prefix can be passed to limit the results to IDs with a common prefix. /// /// ### Arguments @@ -621,8 +728,14 @@ impl PineconeClient { }; let index = Index { - host: endpoint.clone(), - connection: self.new_index_connection(endpoint).await?, + connection: self.new_index_connection(endpoint.clone()).await?, + client: Arc::new(PineconeClient { + openapi_config: Configuration { + base_path: endpoint.clone(), + ..self.openapi_config.clone() + }, + ..self.clone() + }), }; Ok(index) @@ -674,7 +787,7 @@ mod tests { let index = pinecone.index(server.base_url().as_str()).await.unwrap(); - assert_eq!(index.host, server.base_url()); + assert_eq!(index.client.openapi_config.base_path, server.base_url()); } #[tokio::test] @@ -733,4 +846,14 @@ mod tests { .await .expect_err("Expected connection error"); } + + #[tokio::test] + async fn test_serialize_json() { + let json_val = serde_json::json!({ + "_id": "123", + "hello": "world"} + ); + let json = r#"{"_id":"123","hello":"world"}"#; + assert_eq!(json_val.to_string(), json); + } } diff --git a/src/utils/errors.rs b/src/utils/errors.rs index 9cc1735..eac974c 100644 --- a/src/utils/errors.rs +++ b/src/utils/errors.rs @@ -195,8 +195,8 @@ impl From> for PineconeError { } } -// Helper function to handle response errors -fn handle_response_error(source: WrappedResponseContent) -> PineconeError { +/// Helper function to handle response errors +pub fn handle_response_error(source: WrappedResponseContent) -> PineconeError { let status = source.status; let message = source.content.clone(); diff --git a/tests/common.rs b/tests/common.rs index 5128c8f..e29aed3 100644 --- a/tests/common.rs +++ b/tests/common.rs @@ -46,11 +46,6 @@ pub fn get_serverless_index() -> String { std::env::var("SERVERLESS_INDEX_NAME").unwrap() } -/// Returns the name of the pod collection from the environment variable -pub fn get_pod_index() -> String { - std::env::var("POD_INDEX_NAME").unwrap() -} - /// Returns the name of the collection from the environment variable #[allow(dead_code)] pub fn get_collection() -> String { diff --git a/tests/integration_test_control.rs b/tests/integration_test_control.rs index a31b48f..17d6a3c 100644 --- a/tests/integration_test_control.rs +++ b/tests/integration_test_control.rs @@ -1,12 +1,8 @@ -use common::{ - generate_collection_name, generate_index_name, get_collection, get_pod_index, - get_serverless_index, -}; -use pinecone_sdk::models::{Cloud, DeletionProtection, Metric, WaitPolicy}; -use pinecone_sdk::pinecone::{default_client, PineconeClientConfig}; +use common::{generate_index_name, get_serverless_index}; +use pinecone_sdk::models::{Cloud, DeletionProtection, Metric, VectorType, WaitPolicy}; +use pinecone_sdk::pinecone::default_client; use pinecone_sdk::utils::errors::PineconeError; use serial_test::serial; -use std::collections::HashMap; mod common; @@ -52,8 +48,8 @@ async fn test_create_list_indexes() -> Result<(), PineconeError> { "us-west-2", DeletionProtection::Disabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create index"); @@ -67,8 +63,8 @@ async fn test_create_list_indexes() -> Result<(), PineconeError> { "us-west-2", DeletionProtection::Disabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create index"); @@ -88,7 +84,8 @@ async fn test_create_list_indexes() -> Result<(), PineconeError> { assert_eq!(index1.dimension, Some(2)); assert_eq!(index1.metric, Metric::Cosine); let spec1 = index1.spec.serverless.as_ref().unwrap(); - assert_eq!(spec1.cloud, Cloud::Aws); + let spec1_cloud: Cloud = spec1.cloud.into(); + assert_eq!(spec1_cloud, Cloud::Aws); assert_eq!(spec1.region, "us-west-2"); let index2 = indexes @@ -100,7 +97,8 @@ async fn test_create_list_indexes() -> Result<(), PineconeError> { assert_eq!(index2.dimension, Some(2)); assert_eq!(index2.metric, Metric::Dotproduct); let spec2 = index2.spec.serverless.as_ref().unwrap(); - assert_eq!(spec2.cloud, Cloud::Aws); + let spec2_cloud: Cloud = spec2.cloud.into(); + assert_eq!(spec2_cloud, Cloud::Aws); assert_eq!(spec2.region, "us-west-2"); pinecone @@ -131,8 +129,8 @@ async fn test_create_delete_index() -> Result<(), PineconeError> { "us-west-2", DeletionProtection::Disabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create index"); @@ -142,7 +140,8 @@ async fn test_create_delete_index() -> Result<(), PineconeError> { assert_eq!(response.metric, Metric::Euclidean); let spec = response.spec.serverless.unwrap(); - assert_eq!(spec.cloud, Cloud::Aws); + let spec_cloud: Cloud = spec.cloud.into(); + assert_eq!(spec_cloud, Cloud::Aws); assert_eq!(spec.region, "us-west-2"); pinecone @@ -153,98 +152,6 @@ async fn test_create_delete_index() -> Result<(), PineconeError> { Ok(()) } -#[tokio::test] -async fn test_create_pod_index() -> Result<(), PineconeError> { - let pinecone = default_client().expect("Failed to create Pinecone instance"); - - let name = &generate_index_name(); - - let response = pinecone - .create_pod_index( - name, - 2, - Metric::Euclidean, - "us-west1-gcp", - "p1.x1", - 1, - 1, - 1, - DeletionProtection::Disabled, - None, - None, - WaitPolicy::NoWait, - None, - "dense".to_string(), - ) - .await - .expect("Failed to create index"); - - assert_eq!(response.name, name.to_string()); - assert_eq!(response.dimension, Some(2)); - assert_eq!(response.metric, Metric::Euclidean); - - let spec = response.spec.pod.unwrap(); - assert_eq!(spec.environment, "us-west1-gcp"); - assert_eq!(spec.replicas, Some(1)); - assert_eq!(spec.shards, Some(1)); - assert_eq!(spec.pod_type, "p1.x1"); - assert_eq!(spec.pods, Some(1)); - assert_eq!(spec.source_collection, None); - - pinecone - .delete_index(name) - .await - .expect("Failed to delete index"); - - Ok(()) -} - -#[tokio::test] -async fn test_create_pod_index_collection() -> Result<(), PineconeError> { - let pinecone = default_client().expect("Failed to create Pinecone instance"); - - let name = &generate_index_name(); - - let response = pinecone - .create_pod_index( - name, - 12, - Metric::Euclidean, - "us-east-1-aws", - "p1.x1", - 1, - 1, - 1, - DeletionProtection::Disabled, - None, - Some("valid-collection"), - WaitPolicy::NoWait, - None, - "dense".to_string(), - ) - .await - .expect("Failed to create index"); - - assert_eq!(response.name, name.to_string()); - assert_eq!(response.dimension, Some(12)); - assert_eq!(response.metric, Metric::Euclidean); - - let spec = response.spec.pod.unwrap(); - assert_eq!(spec.environment, "us-east-1-aws"); - assert_eq!(spec.replicas, Some(1)); - assert_eq!(spec.shards, Some(1)); - assert_eq!(spec.pod_type, "p1.x1"); - assert_eq!(spec.pods, Some(1)); - assert_eq!(spec.source_collection, Some("valid-collection".to_string())); - - pinecone - .delete_index(name) - .await - .expect("Failed to delete index"); - - Ok(()) -} - #[tokio::test] async fn test_delete_index_err() -> Result<(), PineconeError> { let pinecone = default_client().expect("Failed to create Pinecone instance"); @@ -264,10 +171,10 @@ async fn test_configure_index() -> Result<(), PineconeError> { pinecone .configure_index( - &get_pod_index(), + &get_serverless_index(), Some(DeletionProtection::Enabled), - Some(1), - Some("s1.x1"), + None, + None, None, None, ) @@ -291,8 +198,8 @@ async fn test_configure_deletion_protection() -> Result<(), PineconeError> { "us-east-1", DeletionProtection::Enabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create index"); @@ -322,66 +229,6 @@ async fn test_configure_deletion_protection() -> Result<(), PineconeError> { Ok(()) } -#[tokio::test] -async fn test_configure_optional_deletion_prot() -> Result<(), PineconeError> { - let pinecone = default_client().expect("Failed to create Pinecone instance"); - - let index_name = &generate_index_name(); - pinecone - .create_pod_index( - index_name, - 2, - Metric::Cosine, - "us-east-1-aws", - "p1.x1", - 1, - 1, - 1, - DeletionProtection::Enabled, - None, - None, - WaitPolicy::NoWait, - None, - "dense".to_string(), - ) - .await - .expect("Failed to create index"); - - pinecone - .configure_index(index_name, None, Some(2), None, None, None) - .await - .expect("Failed to configure index"); - - let response = pinecone - .delete_index(index_name) - .await - .expect_err("Expected to fail to delete index"); - - assert!(matches!( - response, - PineconeError::ActionForbiddenError { source: _ } - )); - - pinecone - .configure_index( - index_name, - Some(DeletionProtection::default()), - None, - None, - None, - None, - ) - .await - .expect("Failed to configure index"); - - pinecone - .delete_index(index_name) - .await - .expect("Failed to delete collection"); - - Ok(()) -} - #[tokio::test] async fn test_configure_serverless_index_err() -> Result<(), PineconeError> { let pinecone = default_client().expect("Failed to create Pinecone instance"); @@ -419,141 +266,3 @@ async fn test_configure_invalid_index_err() -> Result<(), PineconeError> { Ok(()) } - -// #[tokio::test] -// #[serial] -// async fn test_create_delete_collection() -> Result<(), PineconeError> { -// let pinecone = default_client().expect("Failed to create Pinecone instance"); -// -// let collection_name = generate_collection_name(); -// -// let index_name = &get_pod_index(); -// loop { -// if match pinecone.describe_index(index_name).await { -// Ok(index) => { -// index.status.ready && (index.status.state == pinecone_sdk::models::State::Ready) -// } -// Err(_) => false, -// } { -// break; -// } -// tokio::time::sleep(Duration::from_millis(1000)).await; -// } -// -// let response = pinecone -// .create_collection(&collection_name, index_name) -// .await -// .expect("Failed to create collection"); -// -// assert_eq!(response.name, collection_name.to_string()); -// -// pinecone -// .delete_collection(&collection_name) -// .await -// .expect("Failed to delete collection"); -// -// Ok(()) -// } - -#[tokio::test] -async fn test_create_collection_serverless_err() -> Result<(), PineconeError> { - let pinecone = default_client().expect("Failed to create Pinecone instance"); - - let collection_name = generate_collection_name(); - - pinecone - .create_collection(&collection_name, &get_serverless_index()) - .await - .expect_err("Expected to fail creating collection from serverless"); - - Ok(()) -} - -#[tokio::test] -async fn test_create_collection_invalid_err() -> Result<(), PineconeError> { - let pinecone = default_client().expect("Failed to create Pinecone instance"); - - let collection_name = generate_collection_name(); - - pinecone - .create_collection(&collection_name, "invalid-index") - .await - .expect_err("Expected to fail creating collection from invalid index"); - - Ok(()) -} - -#[tokio::test] -async fn test_describe_collection() -> Result<(), PineconeError> { - let pinecone = default_client().expect("Failed to create Pinecone instance"); - - let collection_name = get_collection(); - - pinecone - .describe_collection(&collection_name) - .await - .expect("Failed to describe collection"); - - Ok(()) -} - -#[tokio::test] -async fn test_describe_collection_fail() -> Result<(), PineconeError> { - let pinecone = default_client().expect("Failed to create Pinecone instance"); - - pinecone - .describe_collection("invalid-collection") - .await - .expect_err("Expected to fail describing collection"); - - Ok(()) -} - -#[tokio::test] -async fn test_list_collections() -> Result<(), PineconeError> { - let pinecone = default_client().expect("Failed to create Pinecone instance"); - - pinecone - .list_collections() - .await - .expect("Failed to list collections"); - - Ok(()) -} - -#[tokio::test] -async fn test_list_collections_invalid_api_version() -> Result<(), PineconeError> { - let headers: HashMap = [( - pinecone_sdk::pinecone::PINECONE_API_VERSION_KEY.to_string(), - "invalid".to_string(), - )] - .iter() - .cloned() - .collect(); - - let config = PineconeClientConfig { - additional_headers: Some(headers), - ..Default::default() - }; - - let pinecone = config.client().expect("Failed to create client"); - - pinecone - .list_collections() - .await - .expect_err("Expected to fail listing collections due to invalid api version"); - - Ok(()) -} - -#[tokio::test] -async fn test_delete_collection_invalid_collection() -> Result<(), PineconeError> { - let pinecone = default_client().expect("Failed to create Pinecone instance"); - - pinecone - .delete_collection("invalid-collection") - .await - .expect_err("Expected to fail deleting collection"); - - Ok(()) -} diff --git a/tests/integration_test_data.rs b/tests/integration_test_data.rs index ca8015b..b152473 100644 --- a/tests/integration_test_data.rs +++ b/tests/integration_test_data.rs @@ -1,4 +1,4 @@ -use common::{generate_namespace_name, generate_vector, get_pod_index, get_serverless_index}; +use common::{generate_namespace_name, generate_vector, get_serverless_index}; use pinecone_sdk::models::{Kind, Metadata, Namespace, SparseValues, Value, Vector}; use pinecone_sdk::pinecone::default_client; use pinecone_sdk::utils::errors::PineconeError; @@ -112,39 +112,6 @@ async fn test_upsert_sliced_vectors() -> Result<(), PineconeError> { Ok(()) } -#[tokio::test] -async fn test_describe_index_stats_with_filter() -> Result<(), PineconeError> { - let pinecone = default_client().expect("Failed to create Pinecone instance"); - - let host = pinecone - .describe_index(&get_pod_index()) - .await - .unwrap() - .host; - - let mut index = pinecone - .index(host.as_str()) - .await - .expect("Failed to target index"); - - let mut filter = BTreeMap::new(); - filter.insert( - "id".to_string(), - Value { - kind: Some(Kind::BoolValue(false)), - }, - ); - - let describe_index_stats_response = index - .describe_index_stats(Some(Metadata { fields: filter })) - .await - .expect("Failed to describe index stats"); - - assert_eq!(describe_index_stats_response.dimension, 12); - - Ok(()) -} - #[tokio::test] async fn test_describe_index_stats_no_filter() -> Result<(), PineconeError> { let pinecone = default_client().expect("Failed to create Pinecone instance"); @@ -429,79 +396,6 @@ async fn test_delete_all_vectors() -> Result<(), PineconeError> { Ok(()) } -#[tokio::test] -async fn test_delete_by_filter() -> Result<(), PineconeError> { - let pinecone = default_client().expect("Failed to create Pinecone instance"); - - let host = pinecone - .describe_index(&get_pod_index()) - .await - .unwrap() - .host; - - let mut index = pinecone - .index(host.as_str()) - .await - .expect("Failed to target index"); - - let vectors = &[ - Vector { - id: "1".to_string(), - values: vec![1.0; 12], - sparse_values: None, - metadata: Some(Metadata { - fields: vec![( - "key".to_string(), - Value { - kind: Some(Kind::StringValue("value1".to_string())), - }, - )] - .into_iter() - .collect(), - }), - }, - Vector { - id: "2".to_string(), - values: vec![2.0; 12], - sparse_values: None, - metadata: Some(Metadata { - fields: vec![( - "key".to_string(), - Value { - kind: Some(Kind::StringValue("value2".to_string())), - }, - )] - .into_iter() - .collect(), - }), - }, - ]; - - let namespace = &generate_namespace_name(); - index - .upsert(vectors, namespace) - .await - .expect("Failed to upsert"); - - let filter = Metadata { - fields: vec![( - "key".to_string(), - Value { - kind: Some(Kind::StringValue("value1".to_string())), - }, - )] - .into_iter() - .collect(), - }; - - index - .delete_by_filter(filter, namespace) - .await - .expect("Failed to delete all vectors"); - - Ok(()) -} - #[tokio::test] async fn test_fetch_vectors() -> Result<(), PineconeError> { let pinecone = default_client().expect("Failed to create Pinecone instance");