Implement custom types for inference and control plane

src/lib.rs

@@ -18,5 +18,8 @@ pub mod openapi;
 /// Protobuf client for Pinecone.
 pub mod protos;
 
+/// Models for the Pinecone SDK.
+pub mod models;
+
 /// Version information.
 pub mod version;

src/models/embedding.rs

@@ -0,0 +1,21 @@
+use crate::openapi::models::Embedding as OpenApiEmbedding;
+
+/// Embedding
+#[derive(Clone, Default, Debug, PartialEq)]
+pub struct Embedding {
+    /// Embedding values
+    values: Vec<f32>,
+}
+
+impl From<OpenApiEmbedding> for Embedding {
+    fn from(openapi_model: OpenApiEmbedding) -> Self {
+        Embedding {
+            values: openapi_model
+                .values
+                .unwrap_or_default()
+                .into_iter()
+                .map(|x| x as f32)
+                .collect(),
+        }
+    }
+}

src/models/embeddings_list.rs

@@ -0,0 +1,28 @@
+use super::{Embedding, EmbeddingsListUsage};
+use crate::openapi::models::EmbeddingsList as OpenApiEmbeddingsList;
+
+/// EmbeddingsList : Embeddings generated for the input
+#[derive(Clone, Default, Debug, PartialEq)]
+pub struct EmbeddingsList {
+    /// The model used to generate the embeddings.
+    pub model: String,
+    /// The embeddings generated by the model.
+    pub data: Vec<Embedding>,
+    /// The total number of tokens processed.
+    pub usage: EmbeddingsListUsage,
+}
+
+impl From<OpenApiEmbeddingsList> for EmbeddingsList {
+    fn from(openapi_model: OpenApiEmbeddingsList) -> Self {
+        EmbeddingsList {
+            model: openapi_model.model.unwrap_or_default(),
+            data: openapi_model
+                .data
+                .unwrap_or_default()
+                .into_iter()
+                .map(|x| x.into())
+                .collect(),
+            usage: (*openapi_model.usage.unwrap_or_default()).into(),
+        }
+    }
+}

src/models/embeddings_list_usage.rs

@@ -0,0 +1,16 @@
+use crate::openapi::models::EmbeddingsListUsage as OpenApiEmbeddingsListUsage;
+
+/// EmbeddingsListUsage : Usage statistics for model inference including any instruction prefixes
+#[derive(Clone, Default, Debug, PartialEq)]
+pub struct EmbeddingsListUsage {
+    /// The total number of tokens processed.
+    pub total_tokens: i32,
+}
+
+impl From<OpenApiEmbeddingsListUsage> for EmbeddingsListUsage {
+    fn from(openapi_model: OpenApiEmbeddingsListUsage) -> Self {
+        EmbeddingsListUsage {
+            total_tokens: openapi_model.total_tokens.unwrap_or(0),
+        }
+    }
+}

src/models/index_list.rs

@@ -0,0 +1,19 @@
+use super::IndexModel;
+use crate::openapi::models::IndexList as OpenApiIndexList;
+
+/// IndexList : The list of indexes that exist in the project.
+#[derive(Clone, Default, Debug, PartialEq)]
+pub struct IndexList {
+    /// The list of indexes
+    pub indexes: Option<Vec<IndexModel>>,
+}
+
+impl From<OpenApiIndexList> for IndexList {
+    fn from(index_list: OpenApiIndexList) -> Self {
+        IndexList {
+            indexes: index_list
+                .indexes
+                .map(|index| index.into_iter().map(|index| index.into()).collect()),
+        }
+    }
+}

src/models/index_model.rs

@@ -0,0 +1,35 @@
+use super::{DeletionProtection, IndexModelSpec, IndexModelStatus, Metric};
+use crate::openapi::models::index_model::IndexModel as OpenApiIndexModel;
+
+/// IndexModel : The IndexModel describes the configuration and status of a Pinecone index.
+#[derive(Clone, Default, Debug, PartialEq)]
+pub struct IndexModel {
+    /// Index name
+    pub name: String,
+    /// Index dimension
+    pub dimension: i32,
+    /// Index metric
+    pub metric: Metric,
+    /// Index host
+    pub host: String,
+    /// Index deletion protection configuration
+    pub deletion_protection: Option<DeletionProtection>,
+    /// Index specs
+    pub spec: IndexModelSpec,
+    /// Index model specs
+    pub status: IndexModelStatus,
+}
+
+impl From<OpenApiIndexModel> for IndexModel {
+    fn from(openapi_index_model: OpenApiIndexModel) -> Self {
+        IndexModel {
+            name: openapi_index_model.name,
+            dimension: openapi_index_model.dimension,
+            metric: openapi_index_model.metric.into(),
+            host: openapi_index_model.host,
+            deletion_protection: openapi_index_model.deletion_protection,
+            spec: *openapi_index_model.spec,
+            status: *openapi_index_model.status,
+        }
+    }
+}

src/models/metric.rs

@@ -0,0 +1,54 @@
+use crate::openapi::models::create_index_request::Metric as RequestMetric;
+use crate::openapi::models::index_model::Metric as ResponseMetric;
+
+/// The distance metric to be used for similarity search. You can use 'euclidean', 'cosine', or 'dotproduct'.
+#[derive(Clone, Default, Debug, PartialEq)]
+pub enum Metric {
+    /// Cosine similarity
+    #[default]
+    Cosine,
+    /// Euclidean distance similarity
+    Euclidean,
+    /// Dot product similarity
+    Dotproduct,
+}
+
+impl From<RequestMetric> for Metric {
+    fn from(openapi_model: RequestMetric) -> Self {
+        match openapi_model {
+            RequestMetric::Cosine => Metric::Cosine,
+            RequestMetric::Euclidean => Metric::Euclidean,
+            RequestMetric::Dotproduct => Metric::Dotproduct,
+        }
+    }
+}
+
+impl From<ResponseMetric> for Metric {
+    fn from(openapi_model: ResponseMetric) -> Self {
+        match openapi_model {
+            ResponseMetric::Cosine => Metric::Cosine,
+            ResponseMetric::Euclidean => Metric::Euclidean,
+            ResponseMetric::Dotproduct => Metric::Dotproduct,
+        }
+    }
+}
+
+impl From<Metric> for RequestMetric {
+    fn from(model: Metric) -> Self {
+        match model {
+            Metric::Cosine => RequestMetric::Cosine,
+            Metric::Euclidean => RequestMetric::Euclidean,
+            Metric::Dotproduct => RequestMetric::Dotproduct,
+        }
+    }
+}
+
+impl From<Metric> for ResponseMetric {
+    fn from(model: Metric) -> Self {
+        match model {
+            Metric::Cosine => ResponseMetric::Cosine,
+            Metric::Euclidean => ResponseMetric::Euclidean,
+            Metric::Dotproduct => ResponseMetric::Dotproduct,
+        }
+    }
+}

src/models/mod.rs

@@ -0,0 +1,23 @@
+mod embeddings_list;
+pub use self::embeddings_list::EmbeddingsList;
+
+mod embeddings_list_usage;
+pub use self::embeddings_list_usage::EmbeddingsListUsage;
+
+mod metric;
+pub use self::metric::Metric;
+
+mod index_model;
+pub use self::index_model::IndexModel;
+
+mod index_list;
+pub use self::index_list::IndexList;
+
+mod wait_policy;
+pub use self::wait_policy::WaitPolicy;
+
+mod embedding;
+pub use self::embedding::Embedding;
+
+mod openapi;
+pub use self::openapi::*;

src/models/openapi.rs

@@ -0,0 +1,7 @@
+pub use crate::openapi::models::serverless_spec::Cloud;
+pub use crate::openapi::models::{
+    CollectionList, CollectionModel, ConfigureIndexRequest, ConfigureIndexRequestSpec,
+    ConfigureIndexRequestSpecPod, CreateCollectionRequest, DeletionProtection,
+    EmbedRequestParameters, IndexModelSpec, IndexModelStatus, IndexSpec, PodSpec,
+    PodSpecMetadataConfig, ServerlessSpec,
+};

src/models/wait_policy.rs

@@ -0,0 +1,17 @@
+use std::time::Duration;
+
+/// Defines the wait policy for index creation.
+#[derive(Clone, Debug, PartialEq)]
+pub enum WaitPolicy {
+    /// Wait for the index to become ready, up to the specified duration.
+    WaitFor(Duration),
+
+    /// Do not wait for the index to become ready -- return immediately.
+    NoWait,
+}
+
+impl Default for WaitPolicy {
+    fn default() -> Self {
+        WaitPolicy::WaitFor(Duration::from_secs(300))
+    }
+}