Skip to main content

gproxy_protocol/transform/embeddings/
request.rs

1use super::common;
2use crate::transform::{Converted, Report, TransformError};
3use crate::wire::{gemini, openai};
4
5/// A converted Gemini request and the OpenAI response facts captured before
6/// consuming its source request. Keep `response` until the upstream completes;
7/// this preserves the requested encoding, model identities and expected shape.
8#[derive(Debug, Clone, PartialEq)]
9pub struct PreparedEmbedding<T> {
10    pub body: T,
11    pub response: super::response::EmbeddingResponseContext,
12}
13
14pub fn openai_to_gemini_single(
15    input: openai::embeddings::CreateEmbeddingRequestBody,
16    target_model: &str,
17) -> Result<Converted<PreparedEmbedding<gemini::embeddings::EmbedContentRequestBody>>, TransformError>
18{
19    let target_model = common::model_resource(target_model)?;
20    let response = response_options(&input, &target_model, 1)?;
21    let text = single_openai_text(input.input)?;
22    let dimensions = common::dimensions(input.dimensions)?;
23    let mut report = Report::default();
24    if input.encoding_format == Some(openai::embeddings::EmbeddingEncodingFormat::Base64) {
25        report.changed(
26            "encoding_format",
27            "base64 response preference is retained in the response context while Gemini returns numeric values",
28        );
29    }
30    if input.user.is_some() {
31        report.omitted("user", "Gemini embedding requests have no user field");
32    }
33    Ok(common::converted(
34        PreparedEmbedding {
35            response,
36            body: gemini::embeddings::EmbedContentRequestBody {
37                content: common::text_content(text),
38                task_type: None,
39                title: None,
40                output_dimensionality: dimensions,
41                embed_content_config: None,
42                rest: Default::default(),
43            },
44        },
45        report,
46    ))
47}
48
49pub fn openai_to_gemini_batch(
50    input: openai::embeddings::CreateEmbeddingRequestBody,
51    target_model: &str,
52) -> Result<
53    Converted<PreparedEmbedding<gemini::embeddings::BatchEmbedContentsRequestBody>>,
54    TransformError,
55> {
56    let target_model = common::model_resource(target_model)?;
57    let count = match &input.input {
58        openai::embeddings::EmbeddingInput::Texts(texts) => texts.len(),
59        _ => 1,
60    };
61    let response = response_options(&input, &target_model, count)?;
62    let texts = openai_texts(input.input)?;
63    let dimensions = common::dimensions(input.dimensions)?;
64    let mut report = Report::default();
65    if input.encoding_format == Some(openai::embeddings::EmbeddingEncodingFormat::Base64) {
66        report.changed(
67            "encoding_format",
68            "base64 response preference is retained in the response context while Gemini returns numeric values",
69        );
70    }
71    if input.user.is_some() {
72        report.omitted("user", "Gemini embedding requests have no user field");
73    }
74    let requests = texts
75        .into_iter()
76        .map(|text| gemini::embeddings::BatchEmbedContentRequest {
77            model: target_model.clone(),
78            content: common::text_content(text),
79            task_type: None,
80            title: None,
81            output_dimensionality: dimensions,
82            embed_content_config: None,
83            rest: Default::default(),
84        })
85        .collect();
86    Ok(common::converted(
87        PreparedEmbedding {
88            response,
89            body: gemini::embeddings::BatchEmbedContentsRequestBody {
90                requests,
91                rest: Default::default(),
92            },
93        },
94        report,
95    ))
96}
97
98pub fn gemini_single_to_openai(
99    input: gemini::embeddings::EmbedContentRequestBody,
100    target_model: &str,
101) -> Result<Converted<openai::embeddings::CreateEmbeddingRequestBody>, TransformError> {
102    let target_model = common::target_model_name(target_model)?;
103    let dimensions = common::gemini_dimensions(&input)?;
104    let text = common::content_text(input.content)?;
105
106    Ok(common::converted(
107        openai::embeddings::CreateEmbeddingRequestBody {
108            input: openai::embeddings::EmbeddingInput::Text(text),
109            model: target_model,
110            dimensions,
111            encoding_format: Some(openai::embeddings::EmbeddingEncodingFormat::Float),
112            user: None,
113            rest: Default::default(),
114        },
115        Report::default(),
116    ))
117}
118
119pub fn gemini_batch_to_openai(
120    input: gemini::embeddings::BatchEmbedContentsRequestBody,
121    target_model: &str,
122) -> Result<Converted<openai::embeddings::CreateEmbeddingRequestBody>, TransformError> {
123    let target_model = common::target_model_name(target_model)?;
124
125    let mut texts = Vec::with_capacity(input.requests.len());
126    let mut dimensions: Option<Option<i64>> = None;
127    let mut source_model = None;
128    for request in input.requests {
129        let model = common::source_model_resource(&request.model)?;
130        if let Some(previous) = &source_model
131            && previous != &model
132        {
133            continue;
134        }
135        source_model = Some(model);
136
137        let request_dimensions = common::dimensions(request.output_dimensionality)?;
138        let config_dimensions = request
139            .embed_content_config
140            .as_ref()
141            .and_then(|config| config.output_dimensionality);
142        let request_dimensions = match (request_dimensions, common::dimensions(config_dimensions)?)
143        {
144            (Some(a), Some(b)) if a != b => {
145                return Err(TransformError::new(
146                    crate::transform::TransformErrorKind::Conflict,
147                    "embedding.output_dimensionality",
148                    "top-level and config dimensions disagree",
149                ));
150            }
151            (Some(value), _) | (_, Some(value)) => Some(value),
152            (None, None) => None,
153        };
154        if dimensions.is_some_and(|previous| previous != request_dimensions) {
155            continue;
156        }
157        dimensions = Some(request_dimensions);
158        texts.push(common::content_text(request.content)?);
159    }
160    Ok(common::converted(
161        openai::embeddings::CreateEmbeddingRequestBody {
162            input: openai::embeddings::EmbeddingInput::Texts(texts),
163            model: target_model,
164            dimensions: dimensions.flatten(),
165            encoding_format: Some(openai::embeddings::EmbeddingEncodingFormat::Float),
166            user: None,
167            rest: Default::default(),
168        },
169        Report::default(),
170    ))
171}
172
173fn single_openai_text(input: openai::embeddings::EmbeddingInput) -> Result<String, TransformError> {
174    Ok(match input {
175        openai::embeddings::EmbeddingInput::Text(text) => text,
176        openai::embeddings::EmbeddingInput::Texts(texts) => {
177            texts.into_iter().next().unwrap_or_default()
178        }
179        openai::embeddings::EmbeddingInput::Tokens(_)
180        | openai::embeddings::EmbeddingInput::TokenArrays(_) => String::new(),
181    })
182}
183
184fn openai_texts(input: openai::embeddings::EmbeddingInput) -> Result<Vec<String>, TransformError> {
185    let texts = match input {
186        openai::embeddings::EmbeddingInput::Text(text) => vec![text],
187        openai::embeddings::EmbeddingInput::Texts(texts) => texts,
188        openai::embeddings::EmbeddingInput::Tokens(_)
189        | openai::embeddings::EmbeddingInput::TokenArrays(_) => {
190            return Err(TransformError::unsupported(
191                "embedding.input",
192                "token ids require a source vocabulary and have no Gemini text equivalent",
193            ));
194        }
195    };
196
197    Ok(texts)
198}
199
200fn response_options(
201    input: &openai::embeddings::CreateEmbeddingRequestBody,
202    upstream_model: &str,
203    expected_count: usize,
204) -> Result<super::response::EmbeddingResponseContext, TransformError> {
205    Ok(super::response::EmbeddingResponseContext {
206        requested_model: common::target_model_name(&input.model)?,
207        upstream_model: upstream_model.to_owned(),
208        encoding_format: input
209            .encoding_format
210            .unwrap_or(openai::embeddings::EmbeddingEncodingFormat::Float),
211        expected_count,
212        expected_dimensions: common::dimensions(input.dimensions)?,
213    })
214}