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::types::messages::Messages;
8use chat_core::types::response::EmbeddingsResponse;
9
10#[async_trait::async_trait]
11impl EmbeddingsProvider for GeminiClient {
12    async fn embed(&self, messages: &mut Messages) -> Result<EmbeddingsResponse, ChatFailure> {
13        let url = format!(
14            "https://generativelanguage.googleapis.com/v1beta/models/{}:embedContent",
15            self.model_name
16        );
17
18        let request_body =
19            GeminiEmbeddingRequest::from_core(messages, self.embeddings_config.as_ref())
20                .map_err(ChatFailure::from_err)?;
21
22        let res = self
23            .http_client
24            .post(&url)
25            .header("x-goog-api-key", &self.api_key)
26            .json(&request_body)
27            .send()
28            .await
29            .map_err(|e| ChatFailure::from_err(ChatError::Network(e.to_string())))?;
30
31        let res = handle_gemini_error(res).await?;
32
33        let gemini_data: GeminiEmbeddingResponse = res
34            .json()
35            .await
36            .map_err(|e| ChatFailure::from_err(ChatError::InvalidResponse(e.to_string())))?;
37
38        gemini_data
39            .into_core_embeddings_response()
40            .map_err(ChatFailure::from_err)
41    }
42}