initial integrated inf impl

ssmith-pc · ssmith-pc · commit 00d1151868d2 · 2025-04-20T15:37:21.000-07:00
diff --git a/src/models/cloud.rs 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<crate::openapi::models::serverless_spec::Cloud> 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<Cloud> 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<crate::openapi::models::create_index_for_model_request::Cloud> 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<Cloud> 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
@@ -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<HashMap<String, String>>,
+    /// Index embedding configuration
+    pub embed: Option<ModelIndexEmbed>,
+    /// Index vector type
+    pub vector_type: VectorType,
 }
 
 impl From<OpenApiIndexModel> for IndexModel {
@@ -30,6 +39,84 @@ impl From<OpenApiIndexModel> 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<FieldMapEntry>,
+    /// 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<Metric>,
+    /// The desired vector dimension, if supported by the model.
+    pub dimension: Option<i32>,
+    /// The read parameters for the embedding model.
+    pub read_parameters: Option<HashMap<String, ModelParameterValue>>,
+    /// The write parameters for the embedding model.
+    pub write_parameters: Option<HashMap<String, ModelParameterValue>>,
+}
+
+impl From<CreateIndexForModelOptions> 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
@@ -52,3 +52,35 @@ impl From<Metric> for ResponseMetric {
         }
     }
 }
+
+impl From<Metric> 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<crate::openapi::models::create_index_for_model_request_embed::Metric> 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
@@ -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
@@ -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<DeletionProtection>,
+        tags: Option<HashMap<String, String>>,
+        timeout: WaitPolicy,
+    ) -> Result<IndexModel, PineconeError> {
+        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
@@ -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
