portkey_sdk/model/

embeddings.rs

use serde::{Deserialize, Serialize};

/// Input for the embeddings API.
///
/// The input can be a single string, an array of strings, an array of token integers,
/// or an array of token arrays.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum EmbeddingInput {
    /// A single string input
    String(String),

    /// An array of strings
    StringArray(Vec<String>),

    /// An array of token integers
    TokenArray(Vec<i32>),

    /// An array of token arrays
    TokenArrayArray(Vec<Vec<i32>>),
}

/// The format to return embeddings in.
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum EncodingFormat {
    /// Return embeddings as floating point arrays (default)
    #[default]
    Float,

    /// Return embeddings encoded as base64 strings
    Base64,
}

/// Request to create embeddings.
///
/// # Example
///
/// ```rust
/// use portkey_sdk::model::{CreateEmbeddingRequest, EmbeddingInput, EncodingFormat};
///
/// let request = CreateEmbeddingRequest {
///     model: "text-embedding-ada-002".to_string(),
///     input: EmbeddingInput::String("The quick brown fox jumped over the lazy dog".to_string()),
///     encoding_format: Some(EncodingFormat::Float),
///     dimensions: None,
///     user: None,
/// };
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateEmbeddingRequest {
    /// ID of the model to use.
    ///
    /// Common models:
    /// - `text-embedding-ada-002`
    /// - `text-embedding-3-small`
    /// - `text-embedding-3-large`
    pub model: String,

    /// Input text to embed.
    ///
    /// Can be a string, array of strings, array of tokens, or array of token arrays.
    /// The input must not exceed the max input tokens for the model (8192 tokens for
    /// `text-embedding-ada-002`), cannot be an empty string, and any array must be
    /// 2048 dimensions or less.
    pub input: EmbeddingInput,

    /// The format to return the embeddings in.
    ///
    /// Can be either `float` or `base64`. Defaults to `float`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub encoding_format: Option<EncodingFormat>,

    /// The number of dimensions the resulting output embeddings should have.
    ///
    /// Only supported in `text-embedding-3` and later models.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub dimensions: Option<i32>,

    /// A unique identifier representing your end-user.
    ///
    /// This can help providers monitor and detect abuse.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,
}

impl CreateEmbeddingRequest {
    /// Creates a new embedding request with the minimum required fields.
    ///
    /// # Arguments
    ///
    /// * `model` - The model ID to use (e.g., "text-embedding-ada-002", "text-embedding-3-small")
    /// * `input` - The input text or texts to embed
    ///
    /// # Example
    ///
    /// ```
    /// use portkey_sdk::model::{CreateEmbeddingRequest, EmbeddingInput};
    ///
    /// let request = CreateEmbeddingRequest::new(
    ///     "text-embedding-ada-002",
    ///     "Hello, world!",
    /// );
    /// ```
    pub fn new(model: impl Into<String>, input: impl Into<EmbeddingInput>) -> Self {
        Self {
            model: model.into(),
            input: input.into(),
            encoding_format: None,
            dimensions: None,
            user: None,
        }
    }
}

impl From<String> for EmbeddingInput {
    fn from(s: String) -> Self {
        EmbeddingInput::String(s)
    }
}

impl From<&str> for EmbeddingInput {
    fn from(s: &str) -> Self {
        EmbeddingInput::String(s.to_string())
    }
}

impl From<Vec<String>> for EmbeddingInput {
    fn from(v: Vec<String>) -> Self {
        EmbeddingInput::StringArray(v)
    }
}

/// An embedding vector returned by the embedding endpoint.
///
/// # Example
///
/// ```json
/// {
///   "object": "embedding",
///   "embedding": [
///     0.0023064255,
///     -0.009327292,
///     ...
///     -0.0028842222
///   ],
///   "index": 0
/// }
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Embedding {
    /// The index of the embedding in the list of embeddings.
    pub index: i32,

    /// The object type, which is always "embedding".
    pub object: String,

    /// The embedding vector.
    ///
    /// This is a list of floats. The length of the vector depends on the model.
    /// For `text-embedding-ada-002`, this will be 1536 floats.
    pub embedding: Vec<f64>,
}

/// Usage statistics for an embeddings request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingUsage {
    /// The number of tokens used by the prompt.
    pub prompt_tokens: i32,

    /// The total number of tokens used by the request.
    pub total_tokens: i32,
}

/// Response from the embeddings API.
///
/// # Example
///
/// ```rust
/// use portkey_sdk::model::{CreateEmbeddingResponse, Embedding, EmbeddingUsage};
///
/// let response = CreateEmbeddingResponse {
///     object: "list".to_string(),
///     model: "text-embedding-ada-002".to_string(),
///     data: vec![
///         Embedding {
///             index: 0,
///             object: "embedding".to_string(),
///             embedding: vec![0.0023064255, -0.009327292],
///         }
///     ],
///     usage: EmbeddingUsage {
///         prompt_tokens: 8,
///         total_tokens: 8,
///     },
/// };
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateEmbeddingResponse {
    /// The object type, which is always "list".
    pub object: String,

    /// The name of the model used to generate the embedding.
    pub model: String,

    /// The list of embeddings generated by the model.
    pub data: Vec<Embedding>,

    /// The usage information for the request.
    pub usage: EmbeddingUsage,
}