rig_bedrock/
embedding.rs

1use aws_smithy_types::Blob;
2use rig_core::embeddings::{self, Embedding, EmbeddingError};
3use serde::{Deserialize, Serialize};
4
5use crate::{client::Client, types::errors::AwsSdkInvokeModelError};
6
7#[derive(Serialize)]
8#[serde(rename_all = "camelCase")]
9pub struct EmbeddingRequest {
10    pub input_text: String,
11    pub dimensions: usize,
12    pub normalize: bool,
13}
14
15#[derive(Deserialize, Debug)]
16#[serde(rename_all = "camelCase")]
17pub struct EmbeddingResponse {
18    pub embedding: Vec<f64>,
19    pub input_text_token_count: usize,
20}
21
22// The model-id string values are canonically defined in `crate::completion`;
23// these aliases keep this module's historical public names.
24pub use crate::completion::{
25    AMAZON_TITAN_EMBEDDINGS_G1_TEXT as AMAZON_TITAN_EMBED_TEXT_V1,
26    AMAZON_TITAN_MULTIMODAL_EMBEDDINGS_G1 as AMAZON_TITAN_EMBED_IMAGE_V1,
27    AMAZON_TITAN_TEXT_EMBEDDINGS_V2 as AMAZON_TITAN_EMBED_TEXT_V2_0,
28    COHERE_EMBED_ENGLISH as COHERE_EMBED_ENGLISH_V3,
29    COHERE_EMBED_MULTILINGUAL as COHERE_EMBED_MULTILINGUAL_V3,
30};
31
32#[derive(Clone)]
33pub struct EmbeddingModel {
34    client: Client,
35    model: String,
36    ndims: Option<usize>,
37}
38
39impl EmbeddingModel {
40    pub fn new(client: Client, model: impl Into<String>, ndims: Option<usize>) -> Self {
41        Self {
42            client,
43            model: model.into(),
44            ndims,
45        }
46    }
47
48    pub async fn document_to_embeddings(
49        &self,
50        request: EmbeddingRequest,
51    ) -> Result<EmbeddingResponse, EmbeddingError> {
52        let input_document = serde_json::to_string(&request).map_err(EmbeddingError::JsonError)?;
53
54        let model_response = self
55            .client
56            .get_inner()
57            .await
58            .invoke_model()
59            .model_id(self.model.as_str())
60            .content_type("application/json")
61            .accept("application/json")
62            .body(Blob::new(input_document))
63            .send()
64            .await;
65
66        let response = model_response
67            .map_err(|sdk_error| AwsSdkInvokeModelError(sdk_error).into())
68            .map_err(|e: EmbeddingError| e)?;
69
70        let response_str = String::from_utf8(response.body.into_inner())
71            .map_err(|e| EmbeddingError::ResponseError(e.to_string()))?;
72
73        let result: EmbeddingResponse =
74            serde_json::from_str(&response_str).map_err(EmbeddingError::JsonError)?;
75
76        Ok(result)
77    }
78}
79
80impl embeddings::EmbeddingModel for EmbeddingModel {
81    const MAX_DOCUMENTS: usize = 1024;
82
83    type Client = Client;
84
85    fn make(client: &Self::Client, model: impl Into<String>, dims: Option<usize>) -> Self {
86        Self::new(client.clone(), model, dims)
87    }
88
89    fn ndims(&self) -> usize {
90        self.ndims.unwrap_or_default()
91    }
92
93    async fn embed_texts(
94        &self,
95        documents: impl IntoIterator<Item = String> + Send,
96    ) -> Result<Vec<Embedding>, EmbeddingError> {
97        let documents: Vec<String> = documents.into_iter().collect();
98
99        // Deliberately sequential: issuing the requests one at a time keeps
100        // Bedrock's per-account throttling behavior unchanged.
101        let mut results = Vec::new();
102        let mut first_error = None;
103        for doc in documents {
104            let request = EmbeddingRequest {
105                input_text: doc.clone(),
106                dimensions: self.ndims(),
107                normalize: true,
108            };
109            match self.document_to_embeddings(request).await {
110                Ok(embeddings) => results.push(Embedding {
111                    document: doc,
112                    vec: embeddings.embedding,
113                }),
114                Err(err) => {
115                    first_error.get_or_insert(err);
116                }
117            }
118        }
119
120        match first_error {
121            None => Ok(results),
122            Some(err) => Err(EmbeddingError::ResponseError(err.to_string())),
123        }
124    }
125}
rig_bedrock/embedding.rs

rig_bedrock/
embedding.rs
