From f8b5c28c1ada583bc0aa7958c5eeb5cff8b47892 Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Sun, 20 Apr 2025 14:13:42 -0700 Subject: [PATCH 01/16] vector type --- Cargo.lock | 56 +++++++++------------- Cargo.toml | 7 +-- src/models/mod.rs | 3 ++ src/models/vector_type.rs | 13 +++++ src/openapi/models/create_index_request.rs | 3 +- src/openapi/models/index_description.rs | 3 +- src/openapi/models/index_model.rs | 5 +- src/openapi/models/model_index_embed.rs | 3 +- src/pinecone/control.rs | 11 +++-- 9 files changed, 55 insertions(+), 49 deletions(-) create mode 100644 src/models/vector_type.rs 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/models/mod.rs b/src/models/mod.rs index c05b564..f0816e1 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -22,6 +22,9 @@ pub use self::wait_policy::WaitPolicy; mod embedding; pub use self::embedding::Embedding; +mod vector_type; +pub use self::vector_type::VectorType; + pub use crate::openapi::models::{ index_model_status::State, serverless_spec::Cloud, CollectionList, CollectionModel, ConfigureIndexRequest, ConfigureIndexRequestSpec, ConfigureIndexRequestSpecPod, diff --git a/src/models/vector_type.rs b/src/models/vector_type.rs new file mode 100644 index 0000000..edc48f8 --- /dev/null +++ b/src/models/vector_type.rs @@ -0,0 +1,13 @@ +use serde::{Deserialize, Serialize}; +use strum::{Display, EnumString}; + +#[derive( + Debug, Default, Clone, Copy, PartialEq, Eq, Display, EnumString, Serialize, Deserialize, +)] +#[serde(rename_all = "lowercase")] +#[strum(serialize_all = "lowercase")] +pub enum VectorType { + #[default] + Dense, + 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..44cfbce 100644 --- a/src/openapi/models/index_description.rs +++ b/src/openapi/models/index_description.rs @@ -10,6 +10,7 @@ use crate::openapi::models; use serde::{Deserialize, Serialize}; +use crate::models::VectorType; /// IndexDescription : The response for the `describe_index_stats` operation. #[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)] @@ -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..583421c 100644 --- a/src/pinecone/control.rs +++ b/src/pinecone/control.rs @@ -9,7 +9,8 @@ use crate::utils::errors::PineconeError; use crate::models::{ Cloud, CollectionList, CollectionModel, ConfigureIndexRequest, ConfigureIndexRequestSpec, ConfigureIndexRequestSpecPod, CreateCollectionRequest, DeletionProtection, IndexList, - IndexModel, IndexSpec, Metric, PodSpec, PodSpecMetadataConfig, ServerlessSpec, WaitPolicy, + IndexModel, IndexSpec, Metric, PodSpec, PodSpecMetadataConfig, ServerlessSpec, VectorType, + WaitPolicy, }; use crate::openapi::models; @@ -45,7 +46,7 @@ impl PineconeClient { /// Cloud::Aws, // Cloud provider /// "us-east-1", // Region /// DeletionProtection::Enabled, // Deletion protection - /// WaitPolicy::NoWait // Timeout + /// WaitPolicy::NoWait, // Timeout /// None, /// "dense".to_string(), // Vector type /// ).await; @@ -64,7 +65,7 @@ impl PineconeClient { deletion_protection: DeletionProtection, timeout: WaitPolicy, tags: Option>, - vector_type: String, + vector_type: VectorType, ) -> Result { // create request specs let create_index_request_spec = IndexSpec { @@ -167,7 +168,7 @@ impl PineconeClient { source_collection: Option<&str>, timeout: WaitPolicy, tags: Option>, - vector_type: String, + vector_type: VectorType, ) -> Result { // create request specs let indexed = metadata_indexed.map(|i| i.iter().map(|s| s.to_string()).collect()); @@ -194,8 +195,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 From 0250ced89bd5ebe4ac7cd38d40f517572ec74e3d Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Sun, 20 Apr 2025 14:21:48 -0700 Subject: [PATCH 02/16] vector type --- src/models/vector_type.rs | 3 +++ src/openapi/models/index_description.rs | 2 +- src/pinecone/control.rs | 36 ++++++++++++------------- tests/integration_test_control.rs | 16 +++++------ 4 files changed, 30 insertions(+), 27 deletions(-) diff --git a/src/models/vector_type.rs b/src/models/vector_type.rs index edc48f8..a20be6c 100644 --- a/src/models/vector_type.rs +++ b/src/models/vector_type.rs @@ -1,13 +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/index_description.rs b/src/openapi/models/index_description.rs index 44cfbce..51e26ad 100644 --- a/src/openapi/models/index_description.rs +++ b/src/openapi/models/index_description.rs @@ -8,9 +8,9 @@ * Generated by: https://openapi-generator.tech */ +use crate::models::VectorType; use crate::openapi::models; use serde::{Deserialize, Serialize}; -use crate::models::VectorType; /// IndexDescription : The response for the `describe_index_stats` operation. #[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)] diff --git a/src/pinecone/control.rs b/src/pinecone/control.rs index 583421c..19f49d8 100644 --- a/src/pinecone/control.rs +++ b/src/pinecone/control.rs @@ -31,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] @@ -47,8 +47,8 @@ impl PineconeClient { /// "us-east-1", // Region /// DeletionProtection::Enabled, // Deletion protection /// WaitPolicy::NoWait, // Timeout + /// VectorType::Dense, // Vector type /// None, - /// "dense".to_string(), // Vector type /// ).await; /// /// # Ok(()) @@ -64,8 +64,8 @@ impl PineconeClient { region: &str, deletion_protection: DeletionProtection, timeout: WaitPolicy, - tags: Option>, vector_type: VectorType, + tags: Option>, ) -> Result { // create request specs let create_index_request_spec = IndexSpec { @@ -83,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 @@ -120,7 +120,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; /// @@ -145,8 +145,8 @@ impl PineconeClient { /// "imdb_rating"]), /// Some("example-collection"), // Source collection /// WaitPolicy::WaitFor(Duration::from_secs(10)), // Timeout + /// VectorType::Dense, // Vector type /// None, - /// "dense".to_string(), // Vector type /// ) /// .await; /// # Ok(()) @@ -167,8 +167,8 @@ impl PineconeClient { metadata_indexed: Option<&[&str]>, source_collection: Option<&str>, timeout: WaitPolicy, - tags: Option>, vector_type: VectorType, + tags: Option>, ) -> Result { // create request specs let indexed = metadata_indexed.map(|i| i.iter().map(|s| s.to_string()).collect()); @@ -633,8 +633,8 @@ mod tests { "us-east-1", DeletionProtection::Enabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create serverless index"); @@ -697,8 +697,8 @@ mod tests { "us-east-1", DeletionProtection::Enabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create serverless index"); @@ -751,8 +751,8 @@ mod tests { "abc", DeletionProtection::Enabled, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect_err("Expected error when creating serverless index"); @@ -801,8 +801,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"); @@ -851,8 +851,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"); @@ -891,8 +891,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"); @@ -1218,8 +1218,8 @@ mod tests { Some(&["genre", "title", "imdb_rating"]), Some("example-collection"), WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create pod index"); @@ -1304,8 +1304,8 @@ mod tests { None, None, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create pod index"); @@ -1369,8 +1369,8 @@ mod tests { None, Some("example-collection"), WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect_err("Expected create_pod_index to return an error"); @@ -1423,8 +1423,8 @@ mod tests { Some(&["genre", "title", "imdb_rating"]), Some("example-collection"), WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect_err("Expected create_pod_index to return an error"); @@ -1477,8 +1477,8 @@ mod tests { Some(&["genre", "title", "imdb_rating"]), Some("example-collection"), WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect_err("Expected create_pod_index to return an error"); diff --git a/tests/integration_test_control.rs b/tests/integration_test_control.rs index a31b48f..215f681 100644 --- a/tests/integration_test_control.rs +++ b/tests/integration_test_control.rs @@ -2,7 +2,7 @@ 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::models::{Cloud, DeletionProtection, Metric, VectorType, WaitPolicy}; use pinecone_sdk::pinecone::{default_client, PineconeClientConfig}; use pinecone_sdk::utils::errors::PineconeError; use serial_test::serial; @@ -52,8 +52,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 +67,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"); @@ -131,8 +131,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"); @@ -173,8 +173,8 @@ async fn test_create_pod_index() -> Result<(), PineconeError> { None, None, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create index"); @@ -219,8 +219,8 @@ async fn test_create_pod_index_collection() -> Result<(), PineconeError> { None, Some("valid-collection"), WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create index"); @@ -291,8 +291,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"); @@ -341,8 +341,8 @@ async fn test_configure_optional_deletion_prot() -> Result<(), PineconeError> { None, None, WaitPolicy::NoWait, + VectorType::Dense, None, - "dense".to_string(), ) .await .expect("Failed to create index"); From afc434523ab4773566695123430a500242cf8c97 Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Sun, 20 Apr 2025 14:28:56 -0700 Subject: [PATCH 03/16] byoc index --- src/pinecone/control.rs | 91 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 87 insertions(+), 4 deletions(-) diff --git a/src/pinecone/control.rs b/src/pinecone/control.rs index 19f49d8..c0c7072 100644 --- a/src/pinecone/control.rs +++ b/src/pinecone/control.rs @@ -1,8 +1,9 @@ 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, CreateIndexRequest}; use crate::pinecone::PineconeClient; use crate::utils::errors::PineconeError; @@ -65,7 +66,7 @@ impl PineconeClient { deletion_protection: DeletionProtection, timeout: WaitPolicy, vector_type: VectorType, - tags: Option>, + tags: Option>, ) -> Result { // create request specs let create_index_request_spec = IndexSpec { @@ -168,7 +169,7 @@ impl PineconeClient { source_collection: Option<&str>, timeout: WaitPolicy, vector_type: VectorType, - tags: Option>, + tags: Option>, ) -> Result { // create request specs let indexed = metadata_indexed.map(|i| i.iter().map(|s| s.to_string()).collect()); @@ -211,6 +212,88 @@ impl PineconeClient { } } + /// 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' + /// * `deletion_protection: DeletionProtection` - Deletion protection for the index. + /// * `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 + /// * `Result` + /// + /// ### Example + /// ```no_run + /// use pinecone_sdk::models::{IndexModel, Metric, Cloud, WaitPolicy, DeletionProtection, VectorType}; + /// use pinecone_sdk::utils::errors::PineconeError; + /// use std::time::Duration; + /// + /// # #[tokio::main] + /// # async fn main() -> Result<(), PineconeError> { + /// let pinecone = pinecone_sdk::pinecone::default_client()?; + /// + /// // Create a pod index. + /// let response: Result = pinecone.create_byoc_index( + /// "index_name", // Name of the index + /// 10, // Dimension of the index + /// Metric::Cosine, // Distance metric + /// "aws-us-east-1-b921", // Environment + /// DeletionProtection::Enabled, // Deletion protection + /// WaitPolicy::WaitFor(Duration::from_secs(10)), // Timeout + /// VectorType::Dense, // Vector type + /// None, // tags + /// ) + /// .await; + /// # Ok(()) + /// # } + /// ``` + pub async fn create_byoc_index( + &self, + name: &str, + dimension: i32, + metric: Metric, + environment: &str, + deletion_protection: DeletionProtection, + timeout: WaitPolicy, + vector_type: VectorType, + tags: Option>, + ) -> Result { + // create request specs + let spec = ByocSpec { + environment: environment.to_string(), + }; + + let spec = IndexSpec { + serverless: None, + pod: None, + byoc: Some(Box::new(spec)), + }; + + let create_index_request = CreateIndexRequest { + name: name.to_string(), + dimension: Some(dimension), + deletion_protection: Some(deletion_protection), + metric: Some(metric.into()), + spec: Some(Box::new(spec)), + vector_type: Some(vector_type), + tags, + }; + + // make openAPI call + let res = manage_indexes_api::create_index(&self.openapi_config, create_index_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, @@ -363,7 +446,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() { From 00d1151868d2d112dc4c2f656f507451d1b27b9a Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Sun, 20 Apr 2025 15:37:21 -0700 Subject: [PATCH 04/16] initial integrated inf impl --- src/models/cloud.rs | 58 ++++++++++++++++++++ src/models/index_model.rs | 89 ++++++++++++++++++++++++++++++- src/models/metric.rs | 32 +++++++++++ src/models/mod.rs | 15 ++++-- src/pinecone/control.rs | 58 +++++++++++++++++--- tests/integration_test_control.rs | 9 ++-- 6 files changed, 245 insertions(+), 16 deletions(-) create mode 100644 src/models/cloud.rs 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 f0816e1..67369dd 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; @@ -25,11 +27,14 @@ 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, ServerlessSpec, }; pub use crate::protos::{ diff --git a/src/pinecone/control.rs b/src/pinecone/control.rs index c0c7072..efaa1d1 100644 --- a/src/pinecone/control.rs +++ b/src/pinecone/control.rs @@ -3,15 +3,15 @@ use std::collections::HashMap; use std::time::Duration; use crate::openapi::apis::manage_indexes_api; -use crate::openapi::models::{ByocSpec, 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, VectorType, - WaitPolicy, + ConfigureIndexRequestSpecPod, CreateCollectionRequest, CreateIndexForModelOptions, + DeletionProtection, IndexList, IndexModel, IndexSpec, Metric, PodSpec, PodSpecMetadataConfig, + ServerlessSpec, VectorType, WaitPolicy, }; use crate::openapi::models; @@ -71,7 +71,7 @@ impl PineconeClient { // create request specs let create_index_request_spec = IndexSpec { serverless: Some(Box::new(ServerlessSpec { - cloud, + cloud: cloud.into(), region: region.to_string(), })), pod: None, @@ -250,6 +250,7 @@ impl PineconeClient { /// # Ok(()) /// # } /// ``` + #[allow(clippy::too_many_arguments)] pub async fn create_byoc_index( &self, name: &str, @@ -294,6 +295,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, @@ -908,14 +943,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 { @@ -1051,6 +1086,9 @@ mod tests { pod: None, byoc: None, }, + vector_type: VectorType::Dense, + tags: None, + embed: None, }; assert_eq!(index, expected); @@ -1191,6 +1229,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(), @@ -1200,6 +1241,9 @@ mod tests { deletion_protection: None, spec: models::IndexModelSpec::default(), status: models::IndexModelStatus::default(), + embed: None, + tags: None, + vector_type: VectorType::Dense, }, ]), }; diff --git a/tests/integration_test_control.rs b/tests/integration_test_control.rs index 215f681..0e4ede7 100644 --- a/tests/integration_test_control.rs +++ b/tests/integration_test_control.rs @@ -88,7 +88,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 +101,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 @@ -142,7 +144,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 From 1cc3668e4e42d9283404ca957ba8a2e4d235384f Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Mon, 21 Apr 2025 09:00:36 -0700 Subject: [PATCH 05/16] remove pod/collection support --- src/lib.rs | 10 +- src/pinecone/control.rs | 1148 +---------------------------- tests/common.rs | 5 - tests/integration_test_control.rs | 304 +------- tests/integration_test_data.rs | 108 +-- 5 files changed, 45 insertions(+), 1530 deletions(-) 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/pinecone/control.rs b/src/pinecone/control.rs index efaa1d1..b1a98ba 100644 --- a/src/pinecone/control.rs +++ b/src/pinecone/control.rs @@ -8,9 +8,8 @@ use crate::pinecone::PineconeClient; use crate::utils::errors::PineconeError; use crate::models::{ - Cloud, CollectionList, CollectionModel, ConfigureIndexRequest, ConfigureIndexRequestSpec, - ConfigureIndexRequestSpecPod, CreateCollectionRequest, CreateIndexForModelOptions, - DeletionProtection, IndexList, IndexModel, IndexSpec, Metric, PodSpec, PodSpecMetadataConfig, + Cloud, ConfigureIndexRequest, ConfigureIndexRequestSpec, ConfigureIndexRequestSpecPod, + CreateIndexForModelOptions, DeletionProtection, IndexList, IndexModel, IndexSpec, Metric, ServerlessSpec, VectorType, WaitPolicy, }; use crate::openapi::models; @@ -100,118 +99,6 @@ impl PineconeClient { } } - /// Creates a pod 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 - /// * `Result` - /// - /// ### Example - /// ```no_run - /// use pinecone_sdk::models::{IndexModel, Metric, Cloud, WaitPolicy, DeletionProtection, VectorType}; - /// use pinecone_sdk::utils::errors::PineconeError; - /// use std::time::Duration; - /// - /// # #[tokio::main] - /// # async fn main() -> Result<(), PineconeError> { - /// let pinecone = pinecone_sdk::pinecone::default_client()?; - /// - /// // Create a pod index. - /// let response: Result = pinecone.create_pod_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 - /// 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 - /// VectorType::Dense, // Vector type - /// None, - /// ) - /// .await; - /// # Ok(()) - /// # } - /// ``` - #[allow(clippy::too_many_arguments)] - pub async fn create_pod_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, - 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 { - 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, - }; - - let create_index_request = CreateIndexRequest { - name: name.to_string(), - dimension: Some(dimension), - deletion_protection: Some(deletion_protection), - metric: Some(metric.into()), - spec: Some(Box::new(spec)), - vector_type: Some(vector_type), - tags, - }; - - // make openAPI call - let res = manage_indexes_api::create_index(&self.openapi_config, create_index_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), - } - } - /// Creates a BYOC index. /// /// ### Arguments @@ -560,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)] @@ -699,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::*; @@ -1284,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" - } - "#, + }"#, ); }); @@ -1331,367 +1075,42 @@ 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, - VectorType::Dense, - None, - ) - .await - .expect("Failed to create pod index"); - - assert_eq!(create_index_response.name, "index-name"); - assert_eq!(create_index_response.dimension, Some(1536)); - assert_eq!(create_index_response.metric, Metric::Euclidean); - - 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)); + 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_create_pod_index_with_defaults() -> Result<(), PineconeError> { + async fn test_handle_polling_index_err() -> 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": "cosine", - "host": "mock-host", - "spec": { - "pod": { - "environment": "us-east-1-aws", - "pod_type": "p1.x1", - "pods": 1, - "metadata_config": {}, - "replicas": 1, - "shards": 1 + { + "dimension": 1536, + "host": "mock-host", + "metric": "cosine", + "name": "index-name", + "spec": { + "serverless": { + "cloud": "aws", + "region": "us-east-1" + } + }, + "status": { + "ready": false, + "state": "Initializing" } - }, - "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 create_index_response = pinecone - .create_pod_index( - "index-name", - 1536, - Default::default(), - "us-east-1-aws", - "p1.x1", - 1, - 1, - 1, - DeletionProtection::Enabled, - None, - None, - WaitPolicy::NoWait, - VectorType::Dense, - None, - ) - .await - .expect("Failed to create pod 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); - - 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)); - - mock.assert(); - - Ok(()) - } - - #[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, - VectorType::Dense, - None, - ) - .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, - VectorType::Dense, - None, - ) - .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, - VectorType::Dense, - None, - ) - .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" - } - }"#, + }"#, ); }); @@ -2185,501 +1604,4 @@ mod tests { 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", - "dimension": 1536, - "vector_count": 120000, - "environment": "us-east1-gcp" - } - "#, - ); - }); - - // 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()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - // Call create_collection and verify the result - let collection = pinecone - .create_collection("collection1", "index1") - .await - .expect("Failed to create collection"); - - 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); - - mock.assert(); - - Ok(()) - } - - #[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 - } - "#, - ); - }); - - 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"); - - assert!(matches!( - create_collection_response, - PineconeError::CollectionsQuotaExceededError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_create_collection_invalid_name() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(POST).path("/collections"); - then.status(409) - .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 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") - .await - .expect_err("Expected create_collection to return an error"); - - assert!(matches!( - create_collection_response, - PineconeError::InternalServerError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_describe_collection() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(GET).path("/collections/collection-name"); - then.status(200) - .header("content-type", "application/json") - .body( - r#"{ - "dimension": 3, - "environment": "us-east1-gcp", - "name": "tiny-collection", - "size": 3126700, - "status": "Ready", - "vector_count": 99 - }"#, - ); - }); - - // 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()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - // Call describe_collection and verify the result - let collection = pinecone - .describe_collection("collection-name") - .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(), - }; - - assert_eq!(collection, expected); - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_describe_collection_invalid_name() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(GET).path("/collections/invalid-collection"); - then.status(404) - .header("content-type", "application/json") - .body( - r#"{ - "error": "Collection invalid-collection 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 response = pinecone - .describe_collection("invalid-collection") - .await - .expect_err("Expected describe_collection to return an error"); - - assert!(matches!( - response, - PineconeError::CollectionNotFoundError { .. } - )); - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_describe_collection_server_error() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(GET).path("/collections/collection-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 response = pinecone - .describe_collection("collection-name") - .await - .expect_err("Expected describe_collection to return an error"); - - assert!(matches!( - response, - PineconeError::InternalServerError { .. } - )); - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_list_collections() -> 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" - } - ] - }"#, - ); - }); - - // 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()), - ..Default::default() - }; - let pinecone = config.client().expect("Failed to create Pinecone instance"); - - // Call list_collections and verify the result - let collection_list = pinecone - .list_collections() - .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); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_list_collections_error() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(GET).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"); - - // Call list_collections and verify the result - let list_collections_response = pinecone - .list_collections() - .await - .expect_err("Expected to fail to list collections"); - - assert!(matches!( - list_collections_response, - PineconeError::InternalServerError { .. } - )); - 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(); - - Ok(()) - } - - #[tokio::test] - async fn test_delete_collection_not_found() -> 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" - } - "#, - ); - }); - - 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_collection_response = pinecone - .delete_collection("collection-name") - .await - .expect_err("Expected delete_collection to return an error"); - - assert!(matches!( - delete_collection_response, - PineconeError::CollectionNotFoundError { .. } - )); - - mock.assert(); - - Ok(()) - } - - #[tokio::test] - async fn test_delete_collection_internal_error() -> Result<(), PineconeError> { - let server = MockServer::start(); - - let mock = server.mock(|when, then| { - when.method(DELETE).path("/collections/collection-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_collection_response = pinecone - .delete_collection("collection-name") - .await - .expect_err("Expected delete_collection to return an error"); - - assert!(matches!( - delete_collection_response, - PineconeError::InternalServerError { .. } - )); - - mock.assert(); - - Ok(()) - } } 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 0e4ede7..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 common::{generate_index_name, get_serverless_index}; use pinecone_sdk::models::{Cloud, DeletionProtection, Metric, VectorType, WaitPolicy}; -use pinecone_sdk::pinecone::{default_client, PineconeClientConfig}; +use pinecone_sdk::pinecone::default_client; use pinecone_sdk::utils::errors::PineconeError; use serial_test::serial; -use std::collections::HashMap; mod common; @@ -156,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, - VectorType::Dense, - None, - ) - .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, - VectorType::Dense, - None, - ) - .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"); @@ -267,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, ) @@ -325,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, - VectorType::Dense, - None, - ) - .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"); @@ -422,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"); From b63eca0472da11c22c7f45e5e5fd004fef7cef71 Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Mon, 21 Apr 2025 17:33:01 -0700 Subject: [PATCH 06/16] expose inner model --- src/models/mod.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/models/mod.rs b/src/models/mod.rs index 67369dd..cb1046d 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -34,7 +34,8 @@ pub use crate::openapi::models::{ index_model_status::State, CollectionList, CollectionModel, ConfigureIndexRequest, ConfigureIndexRequestSpec, ConfigureIndexRequestSpecPod, CreateCollectionRequest, DeletionProtection, EmbedRequestParameters, IndexModelSpec, IndexModelStatus, IndexSpec, - PodSpec, PodSpecMetadataConfig, ServerlessSpec, + PodSpec, PodSpecMetadataConfig, ServerlessSpec, UpsertRecord, + UpsertResponse as UpsertRecordResponse, }; pub use crate::protos::{ From 3c43c4588360be2523e62a5cbe9caa68d22d0355 Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Mon, 21 Apr 2025 17:51:56 -0700 Subject: [PATCH 07/16] enable reqwest client for data plane --- src/pinecone/data.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/pinecone/data.rs b/src/pinecone/data.rs index ac283d3..1de9d8d 100644 --- a/src/pinecone/data.rs +++ b/src/pinecone/data.rs @@ -2,6 +2,7 @@ use crate::pinecone::PineconeClient; use crate::protos::vector_service_client::VectorServiceClient; use crate::utils::errors::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,7 @@ use crate::models::{ DescribeIndexStatsResponse, FetchResponse, ListResponse, Metadata, Namespace, QueryResponse, SparseValues, UpdateResponse, UpsertResponse, Vector, }; +use crate::openapi::apis::vector_operations_api; use crate::protos; #[derive(Debug, Clone)] @@ -38,6 +40,7 @@ pub struct Index { /// The name of the index. host: String, connection: VectorServiceClient>, + client: Arc, } impl Index { @@ -623,6 +626,7 @@ impl PineconeClient { let index = Index { host: endpoint.clone(), connection: self.new_index_connection(endpoint).await?, + client: Arc::new(self.clone()), }; Ok(index) From 7c7e5fc43e8d9780f3e93f5322274f5d443bc9c2 Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Mon, 21 Apr 2025 22:36:58 -0700 Subject: [PATCH 08/16] add upsert records --- src/pinecone/data.rs | 70 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 69 insertions(+), 1 deletion(-) diff --git a/src/pinecone/data.rs b/src/pinecone/data.rs index 1de9d8d..c3dfa05 100644 --- a/src/pinecone/data.rs +++ b/src/pinecone/data.rs @@ -13,7 +13,8 @@ use crate::models::{ DescribeIndexStatsResponse, FetchResponse, ListResponse, Metadata, Namespace, QueryResponse, SparseValues, UpdateResponse, UpsertResponse, Vector, }; -use crate::openapi::apis::vector_operations_api; +use crate::openapi::apis::vector_operations_api::UpsertRecordsNamespaceError; +use crate::openapi::apis::ResponseContent; use crate::protos; #[derive(Debug, Clone)] @@ -97,6 +98,73 @@ impl Index { Ok(response) } + pub async fn upsert_records( + &mut self, + namespace: &str, + records: Vec, + ) -> Result<(), PineconeError> { + let configuration = self.client.openapi_config.clone(); + + let client = self.client.openapi_config.client.clone(); + + 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.json(&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(ResponseContent { + status, + content, + entity, + } + .into()) + } + } + /// 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 From bae7c97b1051f59b542194e952dd39ca42c6ea0c Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Mon, 21 Apr 2025 23:19:18 -0700 Subject: [PATCH 09/16] fix --- src/pinecone/data.rs | 17 ++++++++++------- src/utils/errors.rs | 4 ++-- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/src/pinecone/data.rs b/src/pinecone/data.rs index c3dfa05..54c9586 100644 --- a/src/pinecone/data.rs +++ b/src/pinecone/data.rs @@ -1,6 +1,6 @@ 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}; @@ -98,6 +98,7 @@ impl Index { Ok(response) } + /// TODO pub async fn upsert_records( &mut self, namespace: &str, @@ -156,12 +157,14 @@ impl Index { Ok(()) } else { let entity: Option = serde_json::from_str(&content).ok(); - Err(ResponseContent { - status, - content, - entity, - } - .into()) + Err(handle_response_error( + ResponseContent { + status, + content, + entity, + } + .into(), + )) } } 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(); From ceddd10972f87f00ea9fffed69e50469bcdc2497 Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Mon, 21 Apr 2025 23:21:36 -0700 Subject: [PATCH 10/16] ref --- src/pinecone/data.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pinecone/data.rs b/src/pinecone/data.rs index 54c9586..b1069ea 100644 --- a/src/pinecone/data.rs +++ b/src/pinecone/data.rs @@ -102,7 +102,7 @@ impl Index { pub async fn upsert_records( &mut self, namespace: &str, - records: Vec, + records: &[serde_json::Value], ) -> Result<(), PineconeError> { let configuration = self.client.openapi_config.clone(); From 0a7bdb2560f397c16a31eefe3110eef3a8298417 Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Mon, 21 Apr 2025 23:37:35 -0700 Subject: [PATCH 11/16] debug --- src/pinecone/data.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/pinecone/data.rs b/src/pinecone/data.rs index b1069ea..b526efd 100644 --- a/src/pinecone/data.rs +++ b/src/pinecone/data.rs @@ -1,4 +1,4 @@ -use crate::pinecone::PineconeClient; +use crate::pinecone::{PineconeClient, PINECONE_API_VERSION_KEY}; use crate::protos::vector_service_client::VectorServiceClient; use crate::utils::errors::{handle_response_error, PineconeError}; use once_cell::sync::Lazy; @@ -141,6 +141,9 @@ impl Index { .map_err(|e| PineconeError::ReqwestError { source: anyhow::Error::from(e), })?; + + println!("xxxxx request {:#?}", req); + let resp = client .execute(req) .await From 056fa7dd0c5ac7a9df068754ec7acc77d6824973 Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Mon, 21 Apr 2025 23:57:34 -0700 Subject: [PATCH 12/16] fix host --- src/pinecone/data.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/src/pinecone/data.rs b/src/pinecone/data.rs index b526efd..2a4a8a5 100644 --- a/src/pinecone/data.rs +++ b/src/pinecone/data.rs @@ -104,13 +104,12 @@ impl Index { namespace: &str, records: &[serde_json::Value], ) -> Result<(), PineconeError> { - let configuration = self.client.openapi_config.clone(); - - let client = self.client.openapi_config.client.clone(); + let configuration = &self.client.openapi_config; + let client = &self.client.openapi_config.client; let uri_str = format!( "{}/records/namespaces/{namespace}/upsert", - configuration.base_path, + self.host, namespace = crate::openapi::apis::urlencode(namespace) ); let mut req_builder = client.request(reqwest::Method::POST, uri_str.as_str()); @@ -142,8 +141,6 @@ impl Index { source: anyhow::Error::from(e), })?; - println!("xxxxx request {:#?}", req); - let resp = client .execute(req) .await From 7995b06cf1969592d9f184e120fd7b53d9fa986e Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Tue, 22 Apr 2025 00:04:17 -0700 Subject: [PATCH 13/16] fix --- src/pinecone/data.rs | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/pinecone/data.rs b/src/pinecone/data.rs index 2a4a8a5..bb833d5 100644 --- a/src/pinecone/data.rs +++ b/src/pinecone/data.rs @@ -1,4 +1,4 @@ -use crate::pinecone::{PineconeClient, PINECONE_API_VERSION_KEY}; +use crate::pinecone::PineconeClient; use crate::protos::vector_service_client::VectorServiceClient; use crate::utils::errors::{handle_response_error, PineconeError}; use once_cell::sync::Lazy; @@ -133,7 +133,7 @@ impl Index { .collect::>() .join("\n"); - req_builder = req_builder.json(&ndjson); + req_builder = req_builder.body(ndjson); let req = req_builder .build() @@ -808,4 +808,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); + } } From 800c5d0c1cac644ef466a030ee0adf2684e4a575 Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Tue, 22 Apr 2025 14:46:19 -0700 Subject: [PATCH 14/16] and basic search wrapper --- src/pinecone/data.rs | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/src/pinecone/data.rs b/src/pinecone/data.rs index bb833d5..fae71ff 100644 --- a/src/pinecone/data.rs +++ b/src/pinecone/data.rs @@ -14,7 +14,11 @@ use crate::models::{ SparseValues, UpdateResponse, UpsertResponse, Vector, }; use crate::openapi::apis::vector_operations_api::UpsertRecordsNamespaceError; -use crate::openapi::apis::ResponseContent; +use crate::openapi::apis::{vector_operations_api, ResponseContent}; +use crate::openapi::models::{ + SearchRecordsRequest, SearchRecordsRequestQuery, SearchRecordsRequestRerank, + SearchRecordsResponse, +}; use crate::protos; #[derive(Debug, Clone)] @@ -168,6 +172,28 @@ impl Index { } } + /// TODO + pub async fn search_records( + &mut self, + namespace: &str, + query: SearchRecordsRequestQuery, + fields: Option>, + rerank: Option, + ) -> Result { + let response = vector_operations_api::search_records_namespace( + &self.client.openapi_config, + namespace, + SearchRecordsRequest { + query: Box::new(query), + fields, + rerank: rerank.map(|r| Box::new(r)), + }, + ) + .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 From f9073262dd069f118f49b7f9e75b0626054256ff Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Tue, 22 Apr 2025 14:59:22 -0700 Subject: [PATCH 15/16] expose types --- src/models/mod.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/models/mod.rs b/src/models/mod.rs index cb1046d..41fed44 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -34,7 +34,8 @@ pub use crate::openapi::models::{ index_model_status::State, CollectionList, CollectionModel, ConfigureIndexRequest, ConfigureIndexRequestSpec, ConfigureIndexRequestSpecPod, CreateCollectionRequest, DeletionProtection, EmbedRequestParameters, IndexModelSpec, IndexModelStatus, IndexSpec, - PodSpec, PodSpecMetadataConfig, ServerlessSpec, UpsertRecord, + PodSpec, PodSpecMetadataConfig, SearchRecordsRequest, SearchRecordsRequestQuery, + SearchRecordsRequestRerank, SearchRecordsResponse, ServerlessSpec, UpsertRecord, UpsertResponse as UpsertRecordResponse, }; From 74f07c1862a8daaea229b069f41abfed3a12ecbf Mon Sep 17 00:00:00 2001 From: Silas Smith <163026730+ssmith-pc@users.noreply.github.com> Date: Tue, 22 Apr 2025 15:33:23 -0700 Subject: [PATCH 16/16] fix client init --- src/pinecone/data.rs | 36 ++++++++++++++++++++++++------------ 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/src/pinecone/data.rs b/src/pinecone/data.rs index fae71ff..728dc2d 100644 --- a/src/pinecone/data.rs +++ b/src/pinecone/data.rs @@ -13,6 +13,7 @@ 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::{ @@ -40,10 +41,7 @@ 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, } @@ -113,7 +111,7 @@ impl Index { let uri_str = format!( "{}/records/namespaces/{namespace}/upsert", - self.host, + configuration.base_path, namespace = crate::openapi::apis::urlencode(namespace) ); let mut req_builder = client.request(reqwest::Method::POST, uri_str.as_str()); @@ -173,10 +171,11 @@ impl Index { } /// TODO - pub async fn search_records( + pub async fn search_records_by_text( &mut self, namespace: &str, - query: SearchRecordsRequestQuery, + query_text: &str, + top_k: u32, fields: Option>, rerank: Option, ) -> Result { @@ -184,9 +183,17 @@ impl Index { &self.client.openapi_config, namespace, SearchRecordsRequest { - query: Box::new(query), + 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(|r| Box::new(r)), + rerank: rerank.map(Box::new), }, ) .await?; @@ -721,9 +728,14 @@ impl PineconeClient { }; let index = Index { - host: endpoint.clone(), - connection: self.new_index_connection(endpoint).await?, - client: Arc::new(self.clone()), + 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) @@ -775,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]