gems/
embed.rs

1use crate::client::Client;
2use crate::messages::Message;
3use crate::models::Model;
4use crate::requests::Content;
5use crate::requests::GeminiEmbedRequest;
6use crate::requests::GeminiEmbedRequests;
7use crate::responses::BatchEmbedContentsResponse;
8use crate::responses::EmbedContentResponse;
9use crate::traits::CTrait;
10use anyhow::Result;
11use derive_builder::Builder;
12use reqwest::Method;
13
14#[derive(Builder, Default, Clone)]
15#[builder(setter(into), default)]
16pub struct Embedding {
17    pub model: Model,
18    pub input: Message,
19}
20
21#[derive(Builder, Default, Clone)]
22#[builder(setter(into), default)]
23pub struct BatchEmbedding {
24    pub model: Model,
25    pub input: Vec<Message>,
26}
27
28#[derive(Clone)]
29pub struct Embeddings {
30    pub client: Client,
31}
32
33impl Embeddings {
34    pub async fn create(&self, params: Embedding) -> Result<EmbedContentResponse> {
35        let request_body = GeminiEmbedRequest {
36            model: format!("models/{}", params.model.to_string()),
37            content: Content {
38                parts: vec![params.input.to_part()],
39            },
40        };
41
42        let req = self
43            .client
44            .request(Method::POST, "embedContent")?
45            .json(&request_body);
46
47        let res = req.send().await?;
48        Ok(res.json().await?)
49    }
50
51    pub async fn batch(&self, params: BatchEmbedding) -> Result<BatchEmbedContentsResponse> {
52        let requests = params
53            .input
54            .into_iter()
55            .map(|message| GeminiEmbedRequest {
56                model: format!("models/{}", params.model.to_string()),
57                content: Content {
58                    parts: vec![message.to_part()],
59                },
60            })
61            .collect::<Vec<GeminiEmbedRequest>>();
62
63        let request_body = GeminiEmbedRequests { requests };
64
65        let req = self
66            .client
67            .request(Method::POST, "batchEmbedContents")?
68            .json(&request_body);
69
70        let res = req.send().await?;
71        Ok(res.json().await?)
72    }
73}
gems/embed.rs

gems/
embed.rs
