-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Implement custom types for inference and control plane (#46)
## Problem Currently, the SDK uses the generated types, however some of these are difficult to use from the user perspective (poorly defined, multiple versions defined, etc). ## Solution We implement a custom `EmbeddingsListUsage`, `EmbeddingsList`, `IndexList`, `IndexModel`, and `Metric` type in a `models` directory. We also export all `openapi` types to the same directory so that the imports will all be consistent. ## Type of Change - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] This change requires a documentation update - [ ] Infrastructure change (CI configs, etc) - [ ] Non-code change (docs, etc) - [ ] None of the above: (explain here) ## Test Plan Test cases should pass. --------- Co-authored-by: Emily Yu <emily.y@pinecone.io>
- Loading branch information
1 parent
eb0400d
commit e511603
Showing
14 changed files
with
305 additions
and
104 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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(), | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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(), | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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), | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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()), | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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, | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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, | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,27 @@ | ||
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; | ||
|
||
pub use crate::openapi::models::{ | ||
serverless_spec::Cloud, CollectionList, CollectionModel, ConfigureIndexRequest, | ||
ConfigureIndexRequestSpec, ConfigureIndexRequestSpecPod, CreateCollectionRequest, | ||
DeletionProtection, EmbedRequestParameters, IndexModelSpec, IndexModelStatus, IndexSpec, | ||
PodSpec, PodSpecMetadataConfig, ServerlessSpec, | ||
}; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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)) | ||
} | ||
} |
Oops, something went wrong.