Skip to main content

gproxy_protocol/transform/openai/embeddings/gemini/
response.rs

1use crate::gemini::embeddings::response::GeminiEmbedContentResponse;
2use crate::openai::embeddings::response::OpenAiEmbeddingsResponse;
3use crate::openai::embeddings::types::{
4    OpenAiCreateEmbeddingResponse, OpenAiEmbeddingData, OpenAiEmbeddingDataObject,
5    OpenAiEmbeddingResponseObject, OpenAiEmbeddingUsage, OpenAiEmbeddingVector,
6};
7use crate::openai::types::OpenAiResponseHeaders;
8use crate::transform::openai::model_list::gemini::utils::{
9    openai_error_response_from_gemini, strip_models_prefix,
10};
11use crate::transform::utils::TransformError;
12
13impl TryFrom<GeminiEmbedContentResponse> for OpenAiEmbeddingsResponse {
14    type Error = TransformError;
15
16    fn try_from(value: GeminiEmbedContentResponse) -> Result<Self, TransformError> {
17        Ok(match value {
18            GeminiEmbedContentResponse::Success {
19                stats_code,
20                headers,
21                body,
22            } => {
23                let model = headers
24                    .extra
25                    .get("x-goog-request-params")
26                    .or_else(|| headers.extra.get("X-Goog-Request-Params"))
27                    .and_then(|params| {
28                        params.split('&').find_map(|pair| {
29                            let (key, value) = pair.split_once('=')?;
30                            if key == "model" {
31                                Some(strip_models_prefix(value))
32                            } else {
33                                None
34                            }
35                        })
36                    })
37                    .unwrap_or_default();
38
39                OpenAiEmbeddingsResponse::Success {
40                    stats_code,
41                    headers: OpenAiResponseHeaders {
42                        extra: headers.extra,
43                    },
44                    body: OpenAiCreateEmbeddingResponse {
45                        data: vec![OpenAiEmbeddingData {
46                            embedding: OpenAiEmbeddingVector::FloatArray(body.embedding.values),
47                            index: 0,
48                            object: OpenAiEmbeddingDataObject::Embedding,
49                        }],
50                        model,
51                        object: OpenAiEmbeddingResponseObject::List,
52                        usage: OpenAiEmbeddingUsage::default(),
53                    },
54                }
55            }
56            GeminiEmbedContentResponse::Error {
57                stats_code,
58                headers,
59                body,
60            } => OpenAiEmbeddingsResponse::Error {
61                stats_code,
62                headers: OpenAiResponseHeaders {
63                    extra: headers.extra,
64                },
65                body: openai_error_response_from_gemini(stats_code, body),
66            },
67        })
68    }
69}