gproxy_protocol/transform/embeddings/
response.rs

1use super::common;
2use crate::transform::{Converted, Report, TransformError};
3use crate::wire::{gemini, openai};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct OpenAiResponseSupplement {
7    pub model: String,
8    pub usage: Option<OpenAiUsageFacts>,
9    pub encoding_format: openai::embeddings::EmbeddingEncodingFormat,
10}
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub struct OpenAiUsageFacts {
14    pub prompt_tokens: i64,
15    pub total_tokens: i64,
16}
17
18/// Request-derived facts kept outside the vendor wire. Response model defaults
19/// to the actual selected upstream; an alias policy belongs to the host.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub struct EmbeddingResponseContext {
22    pub requested_model: String,
23    pub upstream_model: String,
24    pub encoding_format: openai::embeddings::EmbeddingEncodingFormat,
25    pub expected_count: usize,
26    pub expected_dimensions: Option<i64>,
27}
28
29impl EmbeddingResponseContext {
30    pub fn finish_single(
31        &self,
32        input: gemini::embeddings::EmbedContentResponseBody,
33        usage: Option<OpenAiUsageFacts>,
34    ) -> Result<Converted<openai::embeddings::CreateEmbeddingResponseBody>, TransformError> {
35        let _embeddings: Vec<_> = input.embedding.iter().collect();
36
37        gemini_single_response_to_openai(input, &self.supplement(usage))
38    }
39
40    pub fn finish_batch(
41        &self,
42        input: gemini::embeddings::BatchEmbedContentsResponseBody,
43        usage: Option<OpenAiUsageFacts>,
44    ) -> Result<Converted<openai::embeddings::CreateEmbeddingResponseBody>, TransformError> {
45        let _embeddings: Vec<_> = input.embeddings.iter().flatten().collect();
46
47        gemini_batch_response_to_openai(input, &self.supplement(usage))
48    }
49
50    fn supplement(&self, usage: Option<OpenAiUsageFacts>) -> OpenAiResponseSupplement {
51        OpenAiResponseSupplement {
52            model: self.upstream_model.clone(),
53            encoding_format: self.encoding_format,
54            usage,
55        }
56    }
57}
58
59pub fn openai_single_response_to_gemini(
60    input: openai::embeddings::CreateEmbeddingResponseBody,
61) -> Result<Converted<gemini::embeddings::EmbedContentResponseBody>, TransformError> {
62    let embedding = input
63        .data
64        .into_iter()
65        .next()
66        .map(gemini_embedding)
67        .map(crate::transform::optional)
68        .transpose()?
69        .flatten();
70    let report = openai_response_report(&input.usage);
71    Ok(common::converted(
72        gemini::embeddings::EmbedContentResponseBody {
73            embedding,
74            usage_metadata: Some(gemini_usage(input.usage)?),
75            rest: Default::default(),
76        },
77        report,
78    ))
79}
80
81pub fn openai_batch_response_to_gemini(
82    input: openai::embeddings::CreateEmbeddingResponseBody,
83) -> Result<Converted<gemini::embeddings::BatchEmbedContentsResponseBody>, TransformError> {
84    let mut embeddings = Vec::with_capacity(input.data.len());
85    for (expected, item) in input.data.into_iter().enumerate() {
86        let index = i64::try_from(expected).map_err(|_| {
87            TransformError::invalid_result(
88                "embedding.index",
89                "embedding index exceeds integer range",
90            )
91        })?;
92        if item.index != index {
93            return Err(TransformError::invalid_result(
94                "embedding.index",
95                "OpenAI indexes are not in sample order",
96            ));
97        }
98        embeddings.push(gemini_embedding(item)?);
99    }
100    if embeddings.is_empty() {
101        return Err(TransformError::invalid_result(
102            "embedding.data",
103            "batch response has no embeddings",
104        ));
105    }
106    let report = openai_response_report(&input.usage);
107    Ok(common::converted(
108        gemini::embeddings::BatchEmbedContentsResponseBody {
109            embeddings: Some(embeddings),
110            usage_metadata: Some(gemini_usage(input.usage)?),
111            rest: Default::default(),
112        },
113        report,
114    ))
115}
116
117pub fn gemini_single_response_to_openai(
118    input: gemini::embeddings::EmbedContentResponseBody,
119    supplement: &OpenAiResponseSupplement,
120) -> Result<Converted<openai::embeddings::CreateEmbeddingResponseBody>, TransformError> {
121    let embedding = input.embedding.ok_or_else(|| {
122        TransformError::invalid_result("embedding", "Gemini response has no embedding")
123    })?;
124    let item = openai_embedding(0, embedding, supplement.encoding_format)?;
125    let mut report = gemini_response_report(input.usage_metadata.as_ref());
126    let usage = openai_usage(input.usage_metadata, supplement.usage)?;
127    if matches!(
128        supplement.encoding_format,
129        openai::embeddings::EmbeddingEncodingFormat::Base64
130    ) {
131        report.changed(
132            "embedding.encoding_format",
133            "Gemini JSON values are encoded as IEEE754 binary32 little-endian bytes",
134        );
135    }
136    Ok(common::converted(
137        openai::embeddings::CreateEmbeddingResponseBody {
138            data: vec![item],
139            model: nonempty_model(&supplement.model)?,
140            object_: openai::embeddings::EmbeddingListObject::List,
141            usage,
142            rest: Default::default(),
143        },
144        report,
145    ))
146}
147
148pub fn gemini_batch_response_to_openai(
149    input: gemini::embeddings::BatchEmbedContentsResponseBody,
150    supplement: &OpenAiResponseSupplement,
151) -> Result<Converted<openai::embeddings::CreateEmbeddingResponseBody>, TransformError> {
152    let embeddings = input.embeddings.ok_or_else(|| {
153        TransformError::invalid_result("embeddings", "Gemini batch response has no embeddings")
154    })?;
155    if embeddings.is_empty() {
156        return Err(TransformError::invalid_result(
157            "embeddings",
158            "Gemini batch response has no embeddings",
159        ));
160    }
161    let data = embeddings
162        .into_iter()
163        .enumerate()
164        .map(|(index, embedding)| {
165            let index = i64::try_from(index).map_err(|_| {
166                TransformError::invalid_result(
167                    "embedding.index",
168                    "embedding index exceeds integer range",
169                )
170            })?;
171            openai_embedding(index, embedding, supplement.encoding_format)
172        })
173        .filter_map(|value| crate::transform::optional(value).transpose())
174        .collect::<Result<Vec<_>, _>>()?;
175    let mut report = gemini_response_report(input.usage_metadata.as_ref());
176    let usage = openai_usage(input.usage_metadata, supplement.usage)?;
177    if matches!(
178        supplement.encoding_format,
179        openai::embeddings::EmbeddingEncodingFormat::Base64
180    ) {
181        report.changed(
182            "embedding.encoding_format",
183            "Gemini JSON values are encoded as IEEE754 binary32 little-endian bytes",
184        );
185    }
186    Ok(common::converted(
187        openai::embeddings::CreateEmbeddingResponseBody {
188            data,
189            model: nonempty_model(&supplement.model)?,
190            object_: openai::embeddings::EmbeddingListObject::List,
191            usage,
192            rest: Default::default(),
193        },
194        report,
195    ))
196}
197
198fn gemini_embedding(
199    input: openai::embeddings::OpenAiEmbedding,
200) -> Result<gemini::embeddings::ContentEmbedding, TransformError> {
201    let values = common::vector_values(input.embedding)?;
202
203    let dimension = i64::try_from(values.len()).map_err(|_| {
204        TransformError::invalid_result("embedding.dimension", "vector length exceeds integer range")
205    })?;
206    Ok(gemini::embeddings::ContentEmbedding {
207        values: Some(values),
208        shape: Some(vec![dimension]),
209        rest: Default::default(),
210    })
211}
212
213fn openai_embedding(
214    index: i64,
215    input: gemini::embeddings::ContentEmbedding,
216    format: openai::embeddings::EmbeddingEncodingFormat,
217) -> Result<openai::embeddings::OpenAiEmbedding, TransformError> {
218    let values = input
219        .values
220        .ok_or_else(|| TransformError::invalid_result("embedding.values", "values are absent"))?;
221
222    let embedding = match format {
223        openai::embeddings::EmbeddingEncodingFormat::Float => {
224            openai::embeddings::EmbeddingVector::Floats(values)
225        }
226        openai::embeddings::EmbeddingEncodingFormat::Base64 => {
227            openai::embeddings::EmbeddingVector::Base64(common::encode_base64(&values)?)
228        }
229    };
230    Ok(openai::embeddings::OpenAiEmbedding {
231        embedding,
232        index,
233        object_: openai::embeddings::EmbeddingObject::Embedding,
234        rest: Default::default(),
235    })
236}
237
238fn gemini_usage(
239    usage: openai::embeddings::EmbeddingUsage,
240) -> Result<gemini::embeddings::EmbeddingUsageMetadata, TransformError> {
241    if usage.prompt_tokens < 0 || usage.total_tokens < usage.prompt_tokens {
242        return Err(TransformError::invalid_result(
243            "usage",
244            "usage counts must be non-negative and total must cover prompt tokens",
245        ));
246    }
247    Ok(gemini::embeddings::EmbeddingUsageMetadata {
248        prompt_token_count: Some(usage.prompt_tokens),
249        prompt_token_details: None,
250        rest: Default::default(),
251    })
252}
253
254fn openai_usage(
255    source: Option<gemini::embeddings::EmbeddingUsageMetadata>,
256    supplement: Option<OpenAiUsageFacts>,
257) -> Result<openai::embeddings::EmbeddingUsage, TransformError> {
258    let source = source.and_then(|usage| usage.prompt_token_count);
259    let facts = match (source, supplement) {
260        (Some(source), Some(facts)) => {
261            if source != facts.prompt_tokens {
262                return Err(TransformError::new(
263                    crate::transform::TransformErrorKind::Conflict,
264                    "usage.prompt_tokens",
265                    "Gemini usage and supplied usage disagree",
266                ));
267            }
268            OpenAiUsageFacts {
269                prompt_tokens: source,
270                total_tokens: facts.total_tokens,
271            }
272        }
273        (Some(source), None) => OpenAiUsageFacts {
274            prompt_tokens: source,
275            total_tokens: source,
276        },
277        (None, Some(facts)) => facts,
278        (None, None) => {
279            return Err(TransformError::missing_metadata(
280                "usage.prompt_tokens and usage.total_tokens",
281            ));
282        }
283    };
284    if facts.prompt_tokens < 0 || facts.total_tokens < facts.prompt_tokens {
285        return Err(TransformError::invalid_result(
286            "usage",
287            "usage counts must be non-negative and total must cover prompt tokens",
288        ));
289    }
290    Ok(openai::embeddings::EmbeddingUsage {
291        prompt_tokens: facts.prompt_tokens,
292        total_tokens: facts.total_tokens,
293        rest: Default::default(),
294    })
295}
296
297fn nonempty_model(value: &str) -> Result<String, TransformError> {
298    if value.is_empty() {
299        Err(TransformError::missing_metadata("target_model"))
300    } else {
301        Ok(value.to_owned())
302    }
303}
304
305fn openai_response_report(usage: &openai::embeddings::EmbeddingUsage) -> Report {
306    let mut report = Report::default();
307    report.omitted("model", "Gemini embedding responses have no model field");
308    if usage.total_tokens != usage.prompt_tokens {
309        report.omitted(
310            "usage.total_tokens",
311            "Gemini embedding usage has no distinct total-token field",
312        );
313    }
314    report
315}
316
317fn gemini_response_report(usage: Option<&gemini::embeddings::EmbeddingUsageMetadata>) -> Report {
318    let mut report = Report::default();
319    if usage.is_some_and(|usage| usage.prompt_token_details.is_some()) {
320        report.omitted(
321            "usage.prompt_token_details",
322            "OpenAI embedding usage has no per-modality details",
323        );
324    }
325    report
326}
gproxy_protocol/transform/embeddings/response.rs

gproxy_protocol/transform/embeddings/
response.rs
