|
1 | 1 | use std::cmp::min;
|
| 2 | +use std::collections::HashMap; |
2 | 3 | use std::time::Duration;
|
3 | 4 |
|
4 | 5 | use crate::openapi::apis::manage_indexes_api;
|
5 |
| -use crate::openapi::models::CreateIndexRequest; |
| 6 | +use crate::openapi::models::{ByocSpec, CreateIndexRequest}; |
6 | 7 | use crate::pinecone::PineconeClient;
|
7 | 8 | use crate::utils::errors::PineconeError;
|
8 | 9 |
|
@@ -65,7 +66,7 @@ impl PineconeClient {
|
65 | 66 | deletion_protection: DeletionProtection,
|
66 | 67 | timeout: WaitPolicy,
|
67 | 68 | vector_type: VectorType,
|
68 |
| - tags: Option<std::collections::HashMap<String, String>>, |
| 69 | + tags: Option<HashMap<String, String>>, |
69 | 70 | ) -> Result<IndexModel, PineconeError> {
|
70 | 71 | // create request specs
|
71 | 72 | let create_index_request_spec = IndexSpec {
|
@@ -168,7 +169,7 @@ impl PineconeClient {
|
168 | 169 | source_collection: Option<&str>,
|
169 | 170 | timeout: WaitPolicy,
|
170 | 171 | vector_type: VectorType,
|
171 |
| - tags: Option<std::collections::HashMap<String, String>>, |
| 172 | + tags: Option<HashMap<String, String>>, |
172 | 173 | ) -> Result<IndexModel, PineconeError> {
|
173 | 174 | // create request specs
|
174 | 175 | let indexed = metadata_indexed.map(|i| i.iter().map(|s| s.to_string()).collect());
|
@@ -211,6 +212,88 @@ impl PineconeClient {
|
211 | 212 | }
|
212 | 213 | }
|
213 | 214 |
|
| 215 | + /// Creates a BYOC index. |
| 216 | + /// |
| 217 | + /// ### Arguments |
| 218 | + /// * `name: &str` - The name of the index |
| 219 | + /// * `dimension: i32` - The dimension of the index |
| 220 | + /// * `metric: Metric` - The metric to use for the index |
| 221 | + /// * `environment: &str` - The environment where the pod index will be deployed. Example: 'us-east1-gcp' |
| 222 | + /// * `deletion_protection: DeletionProtection` - Deletion protection for the index. |
| 223 | + /// * `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. |
| 224 | + /// |
| 225 | + /// ### Return |
| 226 | + /// * `Result<IndexModel, PineconeError>` |
| 227 | + /// |
| 228 | + /// ### Example |
| 229 | + /// ```no_run |
| 230 | + /// use pinecone_sdk::models::{IndexModel, Metric, Cloud, WaitPolicy, DeletionProtection, VectorType}; |
| 231 | + /// use pinecone_sdk::utils::errors::PineconeError; |
| 232 | + /// use std::time::Duration; |
| 233 | + /// |
| 234 | + /// # #[tokio::main] |
| 235 | + /// # async fn main() -> Result<(), PineconeError> { |
| 236 | + /// let pinecone = pinecone_sdk::pinecone::default_client()?; |
| 237 | + /// |
| 238 | + /// // Create a pod index. |
| 239 | + /// let response: Result<IndexModel, PineconeError> = pinecone.create_byoc_index( |
| 240 | + /// "index_name", // Name of the index |
| 241 | + /// 10, // Dimension of the index |
| 242 | + /// Metric::Cosine, // Distance metric |
| 243 | + /// "aws-us-east-1-b921", // Environment |
| 244 | + /// DeletionProtection::Enabled, // Deletion protection |
| 245 | + /// WaitPolicy::WaitFor(Duration::from_secs(10)), // Timeout |
| 246 | + /// VectorType::Dense, // Vector type |
| 247 | + /// None, // tags |
| 248 | + /// ) |
| 249 | + /// .await; |
| 250 | + /// # Ok(()) |
| 251 | + /// # } |
| 252 | + /// ``` |
| 253 | + pub async fn create_byoc_index( |
| 254 | + &self, |
| 255 | + name: &str, |
| 256 | + dimension: i32, |
| 257 | + metric: Metric, |
| 258 | + environment: &str, |
| 259 | + deletion_protection: DeletionProtection, |
| 260 | + timeout: WaitPolicy, |
| 261 | + vector_type: VectorType, |
| 262 | + tags: Option<HashMap<String, String>>, |
| 263 | + ) -> Result<IndexModel, PineconeError> { |
| 264 | + // create request specs |
| 265 | + let spec = ByocSpec { |
| 266 | + environment: environment.to_string(), |
| 267 | + }; |
| 268 | + |
| 269 | + let spec = IndexSpec { |
| 270 | + serverless: None, |
| 271 | + pod: None, |
| 272 | + byoc: Some(Box::new(spec)), |
| 273 | + }; |
| 274 | + |
| 275 | + let create_index_request = CreateIndexRequest { |
| 276 | + name: name.to_string(), |
| 277 | + dimension: Some(dimension), |
| 278 | + deletion_protection: Some(deletion_protection), |
| 279 | + metric: Some(metric.into()), |
| 280 | + spec: Some(Box::new(spec)), |
| 281 | + vector_type: Some(vector_type), |
| 282 | + tags, |
| 283 | + }; |
| 284 | + |
| 285 | + // make openAPI call |
| 286 | + let res = manage_indexes_api::create_index(&self.openapi_config, create_index_request) |
| 287 | + .await |
| 288 | + .map_err(PineconeError::from)?; |
| 289 | + |
| 290 | + // poll index status |
| 291 | + match self.handle_poll_index(name, timeout).await { |
| 292 | + Ok(_) => Ok(res.into()), |
| 293 | + Err(e) => Err(e), |
| 294 | + } |
| 295 | + } |
| 296 | + |
214 | 297 | // Checks if the index is ready by polling the index status
|
215 | 298 | async fn handle_poll_index(
|
216 | 299 | &self,
|
@@ -363,7 +446,7 @@ impl PineconeClient {
|
363 | 446 | deletion_protection: Option<DeletionProtection>,
|
364 | 447 | replicas: Option<i32>,
|
365 | 448 | pod_type: Option<&str>,
|
366 |
| - tags: Option<std::collections::HashMap<String, String>>, |
| 449 | + tags: Option<HashMap<String, String>>, |
367 | 450 | embed: Option<Box<models::ConfigureIndexRequestEmbed>>,
|
368 | 451 | ) -> Result<IndexModel, PineconeError> {
|
369 | 452 | if replicas.is_none() && pod_type.is_none() && deletion_protection.is_none() {
|
|
0 commit comments