gproxy_protocol/transform/embeddings/
request.rs1use super::common;
2use crate::transform::{Converted, Report, TransformError};
3use crate::wire::{gemini, openai};
4
5#[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}