Skip to content

Commit afc4345

Browse files
committed
byoc index
1 parent 0250ced commit afc4345

File tree

1 file changed

+87
-4
lines changed

1 file changed

+87
-4
lines changed

src/pinecone/control.rs

Lines changed: 87 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
use std::cmp::min;
2+
use std::collections::HashMap;
23
use std::time::Duration;
34

45
use crate::openapi::apis::manage_indexes_api;
5-
use crate::openapi::models::CreateIndexRequest;
6+
use crate::openapi::models::{ByocSpec, CreateIndexRequest};
67
use crate::pinecone::PineconeClient;
78
use crate::utils::errors::PineconeError;
89

@@ -65,7 +66,7 @@ impl PineconeClient {
6566
deletion_protection: DeletionProtection,
6667
timeout: WaitPolicy,
6768
vector_type: VectorType,
68-
tags: Option<std::collections::HashMap<String, String>>,
69+
tags: Option<HashMap<String, String>>,
6970
) -> Result<IndexModel, PineconeError> {
7071
// create request specs
7172
let create_index_request_spec = IndexSpec {
@@ -168,7 +169,7 @@ impl PineconeClient {
168169
source_collection: Option<&str>,
169170
timeout: WaitPolicy,
170171
vector_type: VectorType,
171-
tags: Option<std::collections::HashMap<String, String>>,
172+
tags: Option<HashMap<String, String>>,
172173
) -> Result<IndexModel, PineconeError> {
173174
// create request specs
174175
let indexed = metadata_indexed.map(|i| i.iter().map(|s| s.to_string()).collect());
@@ -211,6 +212,88 @@ impl PineconeClient {
211212
}
212213
}
213214

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+
214297
// Checks if the index is ready by polling the index status
215298
async fn handle_poll_index(
216299
&self,
@@ -363,7 +446,7 @@ impl PineconeClient {
363446
deletion_protection: Option<DeletionProtection>,
364447
replicas: Option<i32>,
365448
pod_type: Option<&str>,
366-
tags: Option<std::collections::HashMap<String, String>>,
449+
tags: Option<HashMap<String, String>>,
367450
embed: Option<Box<models::ConfigureIndexRequestEmbed>>,
368451
) -> Result<IndexModel, PineconeError> {
369452
if replicas.is_none() && pod_type.is_none() && deletion_protection.is_none() {

0 commit comments

Comments
 (0)