Skip to main content

chat_gemini/api/
embedding.rs

1use crate::api::types::error::handle_gemini_error;
2use crate::api::types::request::GeminiEmbeddingRequest;
3use crate::api::types::response::GeminiEmbeddingResponse;
4use crate::client::GeminiClient;
5use chat_core::error::{ChatError, ChatFailure};
6use chat_core::traits::EmbeddingsProvider;
7use chat_core::transport::Transport;
8use chat_core::types::messages::Messages;
9use chat_core::types::response::EmbeddingsResponse;
10
11#[async_trait::async_trait]
12impl<T: Transport> EmbeddingsProvider for GeminiClient<T> {
13    async fn embed(&self, messages: &mut Messages) -> Result<EmbeddingsResponse, ChatFailure> {
14        let path = format!("{}/models/{}:embedContent", self.base_path, self.model_name);
15
16        let request_body =
17            GeminiEmbeddingRequest::from_core(messages, self.embeddings_config.as_ref())
18                .map_err(ChatFailure::from_err)?;
19
20        let body = serde_json::to_vec(&request_body)
21            .map_err(|e| ChatFailure::from_err(ChatError::InvalidResponse(e.to_string())))?;
22
23        let req = chat_core::transport::Request {
24            scheme: self.scheme.clone(),
25            host: self.host.clone(),
26            path,
27            headers: vec![
28                ("x-goog-api-key".into(), self.api_key.clone()),
29                ("Content-Type".into(), "application/json".into()),
30            ],
31            body,
32        };
33
34        let res = self
35            .transport
36            .send(req)
37            .await
38            .map_err(ChatFailure::from_err)?;
39
40        let res = handle_gemini_error(res)?;
41
42        let gemini_data: GeminiEmbeddingResponse = serde_json::from_slice(&res.body)
43            .map_err(|e| ChatFailure::from_err(ChatError::InvalidResponse(e.to_string())))?;
44
45        gemini_data
46            .into_core_embeddings_response()
47            .map_err(ChatFailure::from_err)
48    }
49}