gproxy_protocol/adapt/embeddings/
openai_gemini.rs

1use super::{
2    EmbeddingBatchError, EmbeddingBatchOptions, EmbeddingBatchResult, EmbeddingFailure, at, packing,
3};
4use crate::{
5    WireRequest,
6    adapt::{JsonInvocation, invoke_json},
7    capability::Upstream,
8    codec,
9    transform::{Converted, TransformError, TransformErrorKind, embeddings},
10    wire::{DeclaredFields, gemini::embeddings as g, openai::embeddings as o},
11};
12
13/// Convert an OpenAI embedding batch and call Gemini batchEmbedContents in
14/// source order. `template` is the selected target HTTP endpoint, never the
15/// original client request; its body is replaced and authentication is host-owned.
16pub async fn openai_to_gemini_batch<U: Upstream>(
17    upstream: &U,
18    target: &U::Target,
19    template: WireRequest<()>,
20    input: o::CreateEmbeddingRequestBody,
21    target_model: &str,
22    options: EmbeddingBatchOptions,
23) -> Result<EmbeddingBatchResult<o::CreateEmbeddingResponseBody>, EmbeddingBatchError> {
24    let prepared = embeddings::openai_to_gemini_batch(input, target_model)?;
25    let mut report = prepared.report;
26    let context = prepared.value.response;
27    let batches = packing::gemini_batches(
28        prepared.value.body.requests,
29        &options,
30        upstream.limits().write_bytes,
31    )?;
32    let WireRequest {
33        method,
34        path,
35        query,
36        headers,
37        body: (),
38    } = template;
39    let mut data = Vec::new();
40    let mut usage = o::EmbeddingUsage::builder(0, 0).build();
41    let mut completed = 0;
42    for (index, batch) in batches.into_iter().enumerate() {
43        let expected = batch.requests.len();
44        let request = WireRequest {
45            method: method.clone(),
46            path: path.clone(),
47            query: query.clone(),
48            headers: headers.clone(),
49            body: batch,
50        };
51        let result: JsonInvocation<g::BatchEmbedContentsResponseBody> =
52            invoke_json(upstream, target, request, options.codec)
53                .await
54                .map_err(|error| at(error, completed, data.len()))?;
55        let response = match result {
56            JsonInvocation::Success(response) => {
57                completed += 1;
58                response.body
59            }
60            JsonInvocation::Rejected(response) => {
61                return Err(EmbeddingBatchError {
62                    completed_calls: completed,
63                    completed_items: data.len(),
64                    failure: EmbeddingFailure::Rejected(Box::new(response)),
65                });
66            }
67        };
68        let mut part_context = context.clone();
69        part_context.expected_count = expected;
70        let converted = part_context
71            .finish_batch(
72                response,
73                options.usage_per_call.get(index).copied().flatten(),
74            )
75            .map_err(|error| at(error, completed, data.len()))?;
76        // Consume the response before any same-dialect field reuse.
77        let part = converted.value.into_declared();
78        report.diagnostics.extend(converted.report.diagnostics);
79        let before = data.len();
80        for mut item in part.data {
81            item.index = i64::try_from(data.len()).map_err(|_| {
82                at(
83                    TransformError::new(
84                        TransformErrorKind::Limit,
85                        "embedding.index",
86                        "index exceeds integer range",
87                    ),
88                    completed,
89                    before,
90                )
91            })?;
92            data.push(item);
93        }
94        usage.prompt_tokens = usage
95            .prompt_tokens
96            .checked_add(part.usage.prompt_tokens)
97            .ok_or_else(|| {
98                at(
99                    TransformError::invalid_result(
100                        "usage.prompt_tokens",
101                        "aggregate token count overflow",
102                    ),
103                    completed,
104                    before,
105                )
106            })?;
107        usage.total_tokens = usage
108            .total_tokens
109            .checked_add(part.usage.total_tokens)
110            .ok_or_else(|| {
111                at(
112                    TransformError::invalid_result(
113                        "usage.total_tokens",
114                        "aggregate token count overflow",
115                    ),
116                    completed,
117                    before,
118                )
119            })?;
120        // Bound accumulated output as well as each independently bounded call.
121        codec::encode_json(&data, options.codec)
122            .map_err(|error| at(packing::encoding(error), completed, before))?;
123    }
124    let value = o::CreateEmbeddingResponseBody::builder(
125        data,
126        context.upstream_model,
127        o::EmbeddingListObject::List,
128        usage,
129    )
130    .build();
131    codec::encode_json(&value, options.codec)
132        .map_err(|error| at(packing::encoding(error), completed, value.data.len()))?;
133    Ok(EmbeddingBatchResult {
134        output: Converted { value, report },
135        completed_calls: completed,
136    })
137}
gproxy_protocol/adapt/embeddings/openai_gemini.rs

gproxy_protocol/adapt/embeddings/
openai_gemini.rs
