gproxy_protocol/adapt/embeddings/
gemini_openai.rs

1use super::{
2    EmbeddingBatchError, EmbeddingBatchOptions, EmbeddingBatchResult, EmbeddingFailure, at,
3    openai_packing, packing,
4};
5use crate::{
6    WireRequest,
7    adapt::{JsonInvocation, invoke_json},
8    capability::Upstream,
9    codec,
10    transform::{Converted, Report, TransformError, embeddings},
11    wire::{DeclaredFields, gemini::embeddings as g, openai::embeddings as o},
12};
13
14/// Execute a Gemini embedding batch through the selected OpenAI model. Varying
15/// output dimensions are split into separate calls and returned in input order.
16pub async fn gemini_batch_to_openai<U: Upstream>(
17    upstream: &U,
18    target: &U::Target,
19    template: WireRequest<()>,
20    input: g::BatchEmbedContentsRequestBody,
21    target_model: &str,
22    options: EmbeddingBatchOptions,
23) -> Result<EmbeddingBatchResult<g::BatchEmbedContentsResponseBody>, EmbeddingBatchError> {
24    let input = input.into_declared();
25    if input.requests.is_empty() {
26        return Err(
27            TransformError::shape("embedding.requests", "batch must contain a request").into(),
28        );
29    }
30    let mut prepared = Vec::with_capacity(input.requests.len());
31    let mut report = Report::default();
32    for request in input.requests {
33        if !request.model.starts_with("models/") || request.model.len() == "models/".len() {
34            return Err(TransformError::shape(
35                "embedding.model",
36                "source Gemini model must be a complete models/{name} resource",
37            )
38            .into());
39        }
40        let converted = embeddings::gemini_single_to_openai(
41            g::EmbedContentRequestBody {
42                content: request.content,
43                task_type: request.task_type,
44                title: request.title,
45                output_dimensionality: request.output_dimensionality,
46                embed_content_config: request.embed_content_config,
47                rest: Default::default(),
48            },
49            target_model,
50        )?;
51        report.diagnostics.extend(converted.report.diagnostics);
52        prepared.push(converted.value);
53    }
54    let batches = openai_packing::batches(prepared, &options, upstream.limits().write_bytes)?;
55    let WireRequest {
56        method,
57        path,
58        query,
59        headers,
60        body: (),
61    } = template;
62    let mut vectors = Vec::new();
63    let mut prompt_tokens = 0_i64;
64    let mut completed = 0;
65    for batch in batches {
66        let expected = match &batch.input {
67            o::EmbeddingInput::Texts(values) => values.len(),
68            _ => unreachable!("prepared batch"),
69        };
70        let dimensions = batch.dimensions;
71        let request = WireRequest {
72            method: method.clone(),
73            path: path.clone(),
74            query: query.clone(),
75            headers: headers.clone(),
76            body: batch,
77        };
78        let result: JsonInvocation<o::CreateEmbeddingResponseBody> =
79            invoke_json(upstream, target, request, options.codec)
80                .await
81                .map_err(|error| at(error, completed, vectors.len()))?;
82        let mut response = match result {
83            JsonInvocation::Success(response) => {
84                completed += 1;
85                response.body
86            }
87            JsonInvocation::Rejected(response) => {
88                return Err(EmbeddingBatchError {
89                    completed_calls: completed,
90                    completed_items: vectors.len(),
91                    failure: EmbeddingFailure::Rejected(Box::new(response)),
92                });
93            }
94        };
95        if response.data.len() != expected {
96            return Err(at(
97                TransformError::invalid_result(
98                    "embedding.count",
99                    "response sample count differs from request",
100                ),
101                completed,
102                vectors.len(),
103            ));
104        }
105        // OpenAI indexes define sample order independently of returned array order.
106        // The typed converter below rejects gaps, duplicates and invalid indexes.
107        response.data.sort_unstable_by_key(|item| item.index);
108        let converted = embeddings::openai_batch_response_to_gemini(response)
109            .map_err(|error| at(error, completed, vectors.len()))?;
110        let part = converted.value.into_declared();
111        let values = part.embeddings.ok_or_else(|| {
112            at(
113                TransformError::invalid_result("embeddings", "missing batch vectors"),
114                completed,
115                vectors.len(),
116            )
117        })?;
118        if let Some(expected_dimensions) = dimensions {
119            for value in &values {
120                if value
121                    .values
122                    .as_ref()
123                    .and_then(|v| i64::try_from(v.len()).ok())
124                    != Some(expected_dimensions)
125                {
126                    return Err(at(
127                        TransformError::invalid_result(
128                            "embedding.dimensions",
129                            "response dimensions differ from request",
130                        ),
131                        completed,
132                        vectors.len(),
133                    ));
134                }
135            }
136        }
137        let count = part
138            .usage_metadata
139            .and_then(|v| v.prompt_token_count)
140            .ok_or_else(|| {
141                at(
142                    TransformError::missing_metadata("usage.prompt_tokens"),
143                    completed,
144                    vectors.len(),
145                )
146            })?;
147        prompt_tokens = prompt_tokens.checked_add(count).ok_or_else(|| {
148            at(
149                TransformError::invalid_result(
150                    "usage.prompt_tokens",
151                    "aggregate token count overflow",
152                ),
153                completed,
154                vectors.len(),
155            )
156        })?;
157        let before = vectors.len();
158        vectors.extend(values);
159        report.diagnostics.extend(converted.report.diagnostics);
160        codec::encode_json(&vectors, options.codec)
161            .map_err(|error| at(packing::encoding(error), completed, before))?;
162    }
163    let value = g::BatchEmbedContentsResponseBody::builder()
164        .embeddings(vectors)
165        .usage_metadata(
166            g::EmbeddingUsageMetadata::builder()
167                .prompt_token_count(prompt_tokens)
168                .build(),
169        )
170        .build();
171    codec::encode_json(&value, options.codec).map_err(|error| {
172        at(
173            packing::encoding(error),
174            completed,
175            value.embeddings.as_ref().map_or(0, Vec::len),
176        )
177    })?;
178    Ok(EmbeddingBatchResult {
179        output: Converted { value, report },
180        completed_calls: completed,
181    })
182}
gproxy_protocol/adapt/embeddings/gemini_openai.rs

gproxy_protocol/adapt/embeddings/
gemini_openai.rs
