dynamo_async_openai/types/embedding.rs
1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3//
4// Based on https://github.com/64bit/async-openai/ by Himanshu Neema
5// Original Copyright (c) 2022 Himanshu Neema
6// Licensed under MIT License (see ATTRIBUTIONS-Rust.md)
7//
8// Modifications Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
9// Licensed under Apache 2.0
10
11use base64::engine::{Engine, general_purpose};
12use derive_builder::Builder;
13use serde::{Deserialize, Serialize};
14use utoipa::ToSchema;
15
16use crate::error::OpenAIError;
17
18#[derive(ToSchema, Debug, Serialize, Clone, PartialEq, Deserialize)]
19#[serde(untagged)]
20pub enum EmbeddingInput {
21 String(String),
22 StringArray(Vec<String>),
23 // Minimum value is 0, maximum value is 100257 (inclusive).
24 IntegerArray(Vec<u32>),
25 ArrayOfIntegerArray(Vec<Vec<u32>>),
26}
27
28#[derive(ToSchema, Debug, Serialize, Default, Clone, PartialEq, Deserialize)]
29#[serde(rename_all = "lowercase")]
30pub enum EncodingFormat {
31 #[default]
32 Float,
33 Base64,
34}
35
36#[derive(ToSchema, Debug, Serialize, Default, Clone, Builder, PartialEq, Deserialize)]
37#[builder(name = "CreateEmbeddingRequestArgs")]
38#[builder(pattern = "mutable")]
39#[builder(setter(into, strip_option), default)]
40#[builder(derive(Debug))]
41#[builder(build_fn(error = "OpenAIError"))]
42pub struct CreateEmbeddingRequest {
43 /// ID of the model to use. You can use the
44 /// [List models](https://platform.openai.com/docs/api-reference/models/list)
45 /// API to see all of your available models, or see our
46 /// [Model overview](https://platform.openai.com/docs/models/overview)
47 /// for descriptions of them.
48 pub model: String,
49
50 /// Input text to embed, encoded as a string or array of tokens. To embed multiple inputs in a single request, pass an array of strings 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. [Example Python code](https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken) for counting tokens.
51 pub input: EmbeddingInput,
52
53 /// The format to return the embeddings in. Can be either `float` or [`base64`](https://pypi.org/project/pybase64/). Defaults to float
54 #[serde(skip_serializing_if = "Option::is_none")]
55 pub encoding_format: Option<EncodingFormat>,
56
57 /// A unique identifier representing your end-user, which will help OpenAI
58 /// to monitor and detect abuse. [Learn more](https://platform.openai.com/docs/usage-policies/end-user-ids).
59 #[serde(skip_serializing_if = "Option::is_none")]
60 pub user: Option<String>,
61
62 /// The number of dimensions the resulting output embeddings should have. Only supported in `text-embedding-3` and later models.
63 #[serde(skip_serializing_if = "Option::is_none")]
64 pub dimensions: Option<u32>,
65}
66
67/// Represents an embedding vector returned by embedding endpoint.
68#[derive(ToSchema, Debug, Deserialize, Serialize, Clone, PartialEq)]
69pub struct Embedding {
70 /// The index of the embedding in the list of embeddings.
71 pub index: u32,
72 /// The object type, which is always "embedding".
73 pub object: String,
74 /// The embedding vector, which is a list of floats. The length of vector
75 /// depends on the model as listed in the [embedding guide](https://platform.openai.com/docs/guides/embeddings).
76 pub embedding: Vec<f32>,
77}
78
79#[derive(ToSchema, Debug, Deserialize, Serialize, Clone, PartialEq)]
80pub struct Base64EmbeddingVector(pub String);
81
82impl From<Base64EmbeddingVector> for Vec<f32> {
83 fn from(value: Base64EmbeddingVector) -> Self {
84 let bytes = general_purpose::STANDARD
85 .decode(value.0)
86 .expect("openai base64 encoding to be valid");
87 let chunks = bytes.chunks_exact(4);
88 chunks
89 .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
90 .collect()
91 }
92}
93
94/// Represents an base64-encoded embedding vector returned by embedding endpoint.
95#[derive(ToSchema, Debug, Deserialize, Serialize, Clone, PartialEq)]
96pub struct Base64Embedding {
97 /// The index of the embedding in the list of embeddings.
98 pub index: u32,
99 /// The object type, which is always "embedding".
100 pub object: String,
101 /// The embedding vector, encoded in base64.
102 pub embedding: Base64EmbeddingVector,
103}
104
105#[derive(ToSchema, Debug, Deserialize, Serialize, Clone, PartialEq)]
106pub struct EmbeddingUsage {
107 /// The number of tokens used by the prompt.
108 pub prompt_tokens: u32,
109 /// The total number of tokens used by the request.
110 pub total_tokens: u32,
111}
112
113#[derive(ToSchema, Debug, Deserialize, Clone, PartialEq, Serialize)]
114pub struct CreateEmbeddingResponse {
115 pub object: String,
116 /// The name of the model used to generate the embedding.
117 pub model: String,
118 /// The list of embeddings generated by the model.
119 pub data: Vec<Embedding>,
120 /// The usage information for the request.
121 pub usage: EmbeddingUsage,
122}
123
124#[derive(ToSchema, Debug, Deserialize, Clone, PartialEq, Serialize)]
125pub struct CreateBase64EmbeddingResponse {
126 pub object: String,
127 /// The name of the model used to generate the embedding.
128 pub model: String,
129 /// The list of embeddings generated by the model.
130 pub data: Vec<Base64Embedding>,
131 /// The usage information for the request.
132 pub usage: EmbeddingUsage,
133}