Skip to main content

rig_core/providers/openai/
embedding.rs

1use super::completion::Usage;
2use serde::{Deserialize, Serialize};
3
4/// `text-embedding-3-large` embedding model
5pub const TEXT_EMBEDDING_3_LARGE: &str = "text-embedding-3-large";
6/// `text-embedding-3-small` embedding model
7pub const TEXT_EMBEDDING_3_SMALL: &str = "text-embedding-3-small";
8/// `text-embedding-ada-002` embedding model
9pub const TEXT_EMBEDDING_ADA_002: &str = "text-embedding-ada-002";
10
11#[derive(Debug, Deserialize)]
12pub struct EmbeddingResponse {
13    pub object: String,
14    pub data: Vec<EmbeddingData>,
15    pub model: String,
16    pub usage: Usage,
17}
18
19/// Typed raw response from [`Embeddings`](super::wire::Embeddings).
20/// Missing object and model fields default to empty strings; usage is optional.
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct CompatibleEmbeddingResponse {
23    #[serde(default)]
24    pub object: String,
25    pub data: Vec<EmbeddingData>,
26    #[serde(default)]
27    pub model: String,
28    #[serde(default)]
29    pub usage: Option<Usage>,
30}
31
32#[derive(Debug, Deserialize, Clone, Copy, PartialEq, Eq, Serialize)]
33#[serde(rename_all = "snake_case")]
34pub enum EncodingFormat {
35    Float,
36    Base64,
37}
38
39/// One embedded input. Missing object and index fields default to empty and zero.
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct EmbeddingData {
42    #[serde(default)]
43    pub object: String,
44    pub embedding: Vec<serde_json::Number>,
45    #[serde(default)]
46    pub index: usize,
47}
48
49/// Return default dimensions for a known model identifier, or `None`.
50pub(crate) fn model_dimensions_from_identifier(identifier: &str) -> Option<usize> {
51    match identifier {
52        TEXT_EMBEDDING_3_LARGE => Some(3_072),
53        TEXT_EMBEDDING_3_SMALL | TEXT_EMBEDDING_ADA_002 => Some(1_536),
54        _ => None,
55    }
56}
57
58#[cfg(test)]
59mod tests;