rig_core/providers/gemini/
embedding.rs

1//! Batch text embeddings through the [Gemini API](https://ai.google.dev/api/embeddings).
2//!
3//! ```no_run
4//! use rig_core::providers::gemini::{Gemini, embedding::EMBEDDING_001};
5//!
6//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
7//! let model = Gemini::from_env()?.embedding(EMBEDDING_001, None);
8//! # Ok(())
9//! # }
10//! ```
11
12use serde_json::{Map, Value, json};
13
14use crate::embeddings;
15use crate::error::EncodeError;
16use crate::providers::internal::wire::classify_marker_keyed_frame;
17use crate::wire::{Body, Capabilities, Descriptor, DocumentDecoder, Encoded, Framing, Mode, Wire};
18
19/// `gemini-embedding-001` embedding model (3072 dimensions by default)
20pub const EMBEDDING_001: &str = "gemini-embedding-001";
21/// `text-embedding-004` embedding model (768 dimensions by default)
22pub const EMBEDDING_004: &str = "text-embedding-004";
23
24/// The width the Gemini API embeds `model` at when no width is asked for:
25/// 3072 for [`EMBEDDING_001`], and 768 for [`EMBEDDING_004`] and any model
26/// rig does not know. Every Gemini embedding wire defaults to it.
27///
28/// See <https://ai.google.dev/gemini-api/docs/models#gemini-embedding>
29pub fn default_ndims(model: &str) -> usize {
30    match model {
31        EMBEDDING_001 => 3072,
32        _ => 768,
33    }
34}
35
36/// Gemini's batch embedding endpoint.
37///
38/// `POST /v1beta/models/{model}:batchEmbedContents`, authenticated by the
39/// `key` query parameter the GenerateContent family uses.
40#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
41pub struct Embeddings {
42    /// The provider this wire speaks to.
43    pub provider: super::GeminiConfig,
44    /// The embedding model, as the path names it.
45    pub model: String,
46    /// The `output_dimensionality` every document in the batch asks for.
47    pub ndims: usize,
48}
49
50impl Embeddings {
51    /// Build a wire for `model` with the requested output dimensions, or
52    /// else [`default_ndims`].
53    pub fn new(
54        provider: super::GeminiConfig,
55        model: impl Into<String>,
56        ndims: Option<usize>,
57    ) -> Self {
58        let model = model.into();
59        let ndims = ndims.unwrap_or_else(|| default_ndims(&model));
60        Self {
61            provider,
62            model,
63            ndims,
64        }
65    }
66}
67
68impl Wire for Embeddings {
69    type Op = crate::operation::Embedding;
70    type Payload = crate::wire::Encoded;
71    type Frame = crate::wire::WireFrame;
72    type Event = gemini_api_types::EmbeddingResponse;
73    type Decoder<'id> = DocumentDecoder<Self::Op, Self::Event>;
74    type Reassembler = crate::wire::document::Unreassembled;
75
76    fn describe(&self) -> Descriptor<'_> {
77        Descriptor::new(super::PROVIDER_NAME)
78            .model(self.model.as_str())
79            .capabilities(Capabilities::embedding(1024, self.ndims))
80    }
81
82    fn encode(&self, request: Vec<String>, _mode: Mode) -> Result<Encoded, EncodeError> {
83        let model = format!("models/{}", self.model);
84        let requests = request.iter().map(|doc| {
85            json!({
86                "model": model,
87                "content": { "parts": [{ "text": doc }] },
88                "output_dimensionality": self.ndims,
89            })
90        });
91        let body = Value::Object(Map::from_iter([(
92            "requests".to_owned(),
93            Value::Array(requests.collect()),
94        )]));
95
96        // Avoid allocating formatted batch JSON unless tracing consumes it.
97        if tracing::enabled!(target: "rig::embedding", tracing::Level::TRACE)
98            && let Ok(pretty_body) = serde_json::to_string_pretty(&body)
99        {
100            tracing::trace!(
101                target: "rig::embedding",
102                "Sending embedding request to Gemini API {pretty_body}"
103            );
104        }
105
106        let request = http::Request::post(format!(
107            "{}/v1beta/models/{}:batchEmbedContents?key={}",
108            self.provider.base_url,
109            self.model,
110            self.provider.api_key.expose()
111        ))
112        .header(http::header::CONTENT_TYPE, "application/json")
113        .body(Body::Bytes(serde_json::to_vec(&body)?))?;
114        // `batchEmbedContents` has no streaming variant, so a streamed call
115        // sends the same bytes and reads the same whole reply.
116        Ok(Encoded::new(request, Framing::Whole))
117    }
118
119    fn decoder<'id>(&self) -> Self::Decoder<'id> {
120        DocumentDecoder::new(
121            |body| classify_marker_keyed_frame(body, &["embeddings"]),
122            |reply| {
123                let vectors = reply
124                    .embeddings
125                    .into_iter()
126                    .map(|embedding| embedding.values);
127                // Gemini supplies no usage or response id; the driver attaches the raw body.
128                Ok(embeddings::EmbeddingResponse::from_vectors(vectors))
129            },
130        )
131    }
132}
133
134/// Request and response types for [Gemini embeddings](https://ai.google.dev/api/embeddings).
135///
136/// ```
137/// use rig_core::providers::gemini::embedding::gemini_api_types::EmbeddingValues;
138///
139/// let vector = EmbeddingValues { values: vec![1.into(), 2.into()] };
140/// ```
141pub mod gemini_api_types {
142    use serde::{Deserialize, Serialize};
143
144    #[derive(Debug, Clone, Serialize, Deserialize)]
145    pub struct EmbeddingResponse {
146        pub embeddings: Vec<EmbeddingValues>,
147    }
148
149    #[derive(Debug, Clone, Serialize, Deserialize)]
150    pub struct EmbeddingValues {
151        #[serde(default)]
152        pub values: Vec<f64>,
153    }
154}
155
156impl super::GeminiConfig {
157    /// The `batchEmbedContents` embedding wire. `ndims` defaults from the
158    /// model identifier.
159    pub(crate) fn embedding(&self, model: impl Into<String>, ndims: Option<usize>) -> Embeddings {
160        Embeddings::new(self.clone(), model, ndims)
161    }
162}
163
164#[cfg(test)]
165mod tests;
rig_core/providers/gemini/embedding.rs

rig_core/providers/gemini/
embedding.rs
