Skip to content

Commit

Permalink
Implement custom types for inference and control plane (#46)
Browse files Browse the repository at this point in the history
## 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
ericapywang and emily-emily authored Jul 31, 2024
1 parent eb0400d commit e511603
Show file tree
Hide file tree
Showing 14 changed files with 305 additions and 104 deletions.
3 changes: 3 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
21 changes: 21 additions & 0 deletions src/models/embedding.rs
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(),
}
}
}
28 changes: 28 additions & 0 deletions src/models/embeddings_list.rs
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(),
}
}
}
16 changes: 16 additions & 0 deletions src/models/embeddings_list_usage.rs
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),
}
}
}
19 changes: 19 additions & 0 deletions src/models/index_list.rs
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()),
}
}
}
35 changes: 35 additions & 0 deletions src/models/index_model.rs
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,
}
}
}
54 changes: 54 additions & 0 deletions src/models/metric.rs
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,
}
}
}
27 changes: 27 additions & 0 deletions src/models/mod.rs
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,
};
17 changes: 17 additions & 0 deletions src/models/wait_policy.rs
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))
}
}
Loading

0 comments on commit e511603

Please sign in to comment.