systemprompt_models/ai/
embedding.rs

1//! Provider-neutral text embedding request and response types.
2//!
3//! An [`EmbeddingRequest`] names a model and the texts to embed; the provider
4//! answers with one vector per input, in input order, plus the token usage it
5//! billed. The shapes are the intersection of what the `OpenAI` `/embeddings`,
6//! Vertex AI `:predict` and Gemini `:batchEmbedContents` endpoints return, so
7//! a caller never sees a vendor body. A provider that reports no usage has
8//! its input counted by the caller and the usage marked `estimated`, so a
9//! priced call is never settled at zero.
10//!
11//! Copyright (c) systemprompt.io — Business Source License 1.1.
12//! See <https://systemprompt.io> for licensing details.
13
14use serde::{Deserialize, Serialize};
15use systemprompt_identifiers::{ModelId, ProviderId};
16
17#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
18pub struct EmbeddingRequest {
19    pub model: ModelId,
20    pub inputs: Vec<String>,
21    #[serde(default, skip_serializing_if = "Option::is_none")]
22    pub dimensions: Option<u32>,
23    #[serde(default, skip_serializing_if = "Option::is_none")]
24    pub provider: Option<ProviderId>,
25}
26
27impl EmbeddingRequest {
28    pub fn new(model: impl Into<ModelId>, inputs: Vec<String>) -> Self {
29        Self {
30            model: model.into(),
31            inputs,
32            dimensions: None,
33            provider: None,
34        }
35    }
36
37    #[must_use]
38    pub const fn with_dimensions(mut self, dimensions: u32) -> Self {
39        self.dimensions = Some(dimensions);
40        self
41    }
42
43    #[must_use]
44    pub fn with_provider(mut self, provider: impl Into<ProviderId>) -> Self {
45        self.provider = Some(provider.into());
46        self
47    }
48
49    #[must_use]
50    pub fn provider(&self) -> Option<&ProviderId> {
51        self.provider.as_ref().filter(|p| !p.as_str().is_empty())
52    }
53
54    #[must_use]
55    pub fn with_inputs(&self, inputs: Vec<String>) -> Self {
56        Self {
57            model: self.model.clone(),
58            inputs,
59            dimensions: self.dimensions,
60            provider: self.provider.clone(),
61        }
62    }
63}
64
65#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
66pub struct EmbeddingUsage {
67    pub input_tokens: u32,
68    pub total_tokens: u32,
69    #[serde(default)]
70    pub estimated: bool,
71}
72
73impl EmbeddingUsage {
74    #[must_use]
75    pub const fn reported(input_tokens: u32, total_tokens: u32) -> Self {
76        Self {
77            input_tokens,
78            total_tokens,
79            estimated: false,
80        }
81    }
82
83    #[must_use]
84    pub const fn from_estimate(input_tokens: u32) -> Self {
85        Self {
86            input_tokens,
87            total_tokens: input_tokens,
88            estimated: true,
89        }
90    }
91
92    #[must_use]
93    pub const fn summed_with(self, other: Self) -> Self {
94        Self {
95            input_tokens: self.input_tokens.saturating_add(other.input_tokens),
96            total_tokens: self.total_tokens.saturating_add(other.total_tokens),
97            estimated: self.estimated || other.estimated,
98        }
99    }
100}
101
102#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
103pub struct EmbeddingResponse {
104    pub model: ModelId,
105    pub vectors: Vec<Vec<f32>>,
106    #[serde(default)]
107    pub usage: EmbeddingUsage,
108}
109
110impl EmbeddingResponse {
111    #[must_use]
112    pub const fn empty(model: ModelId) -> Self {
113        Self {
114            model,
115            vectors: Vec::new(),
116            usage: EmbeddingUsage {
117                input_tokens: 0,
118                total_tokens: 0,
119                estimated: false,
120            },
121        }
122    }
123
124    #[must_use]
125    pub fn dimensions(&self) -> Option<usize> {
126        self.vectors.first().map(Vec::len)
127    }
128}
systemprompt_models/ai/embedding.rs

systemprompt_models/ai/
embedding.rs
