rig_core/providers/together/
embedding.rs

1// ================================================================
2//! Together AI Embeddings Integration
3//! From [Together AI Reference](https://docs.together.ai/docs/embeddings-overview)
4// ================================================================
5
6use serde::Deserialize;
7use serde_json::json;
8
9use crate::{
10    embeddings::{self, EmbeddingError},
11    http_client::{self, HttpClientExt},
12};
13
14use super::{Client, client::together_ai_api_types::ApiResponse};
15
16// ================================================================
17// Together AI Embedding API
18// ================================================================
19pub const BGE_BASE_EN_V1_5: &str = "BAAI/bge-base-en-v1.5";
20pub const BGE_LARGE_EN_V1_5: &str = "BAAI/bge-large-en-v1.5";
21pub const BERT_BASE_UNCASED: &str = "bert-base-uncased";
22pub const M2_BERT_2K_RETRIEVAL_ENCODER_V1: &str = "hazyresearch/M2-BERT-2k-Retrieval-Encoder-V1";
23pub const M2_BERT_80M_32K_RETRIEVAL: &str = "togethercomputer/m2-bert-80M-32k-retrieval";
24pub const M2_BERT_80M_2K_RETRIEVAL: &str = "togethercomputer/m2-bert-80M-2k-retrieval";
25pub const M2_BERT_80M_8K_RETRIEVAL: &str = "togethercomputer/m2-bert-80M-8k-retrieval";
26pub const SENTENCE_BERT: &str = "sentence-transformers/msmarco-bert-base-dot-v5";
27pub const UAE_LARGE_V1: &str = "WhereIsAI/UAE-Large-V1";
28
29#[derive(Debug, Deserialize)]
30pub struct EmbeddingResponse {
31    pub model: String,
32    pub object: String,
33    pub data: Vec<EmbeddingData>,
34}
35
36#[derive(Debug, Deserialize)]
37pub struct EmbeddingData {
38    pub object: String,
39    pub embedding: Vec<serde_json::Number>,
40    pub index: usize,
41}
42
43#[derive(Debug, Deserialize)]
44pub struct Usage {
45    pub prompt_tokens: usize,
46    pub total_tokens: usize,
47}
48
49#[derive(Clone)]
50pub struct EmbeddingModel<T = reqwest::Client> {
51    client: Client<T>,
52    pub model: String,
53    ndims: usize,
54}
55
56impl<T> embeddings::EmbeddingModel for EmbeddingModel<T>
57where
58    T: HttpClientExt + Default + Clone + Send + 'static,
59{
60    const MAX_DOCUMENTS: usize = 1024; // This might need to be adjusted based on Together AI's actual limit
61
62    type Client = Client<T>;
63
64    fn make(client: &Self::Client, model: impl Into<String>, dims: Option<usize>) -> Self {
65        Self::new(client.clone(), model, dims.unwrap_or_default())
66    }
67
68    fn ndims(&self) -> usize {
69        self.ndims
70    }
71
72    async fn embed_texts(
73        &self,
74        documents: impl IntoIterator<Item = String>,
75    ) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
76        let documents = documents.into_iter().collect::<Vec<_>>();
77
78        let body = serde_json::to_vec(&json!({
79            "model": self.model,
80            "input": documents,
81        }))?;
82
83        let req = self
84            .client
85            .post("/v1/embeddings")?
86            .body(body)
87            .map_err(|e| EmbeddingError::HttpError(e.into()))?;
88
89        let response = self.client.send(req).await?;
90
91        let status = response.status();
92        if status.is_success() {
93            let response_body: Vec<u8> = response.into_body().await?;
94            let parsed: ApiResponse<EmbeddingResponse> = serde_json::from_slice(&response_body)?;
95
96            match parsed {
97                ApiResponse::Ok(response) => {
98                    if response.data.len() != documents.len() {
99                        return Err(EmbeddingError::ResponseError(
100                            "Response data length does not match input length".into(),
101                        ));
102                    }
103
104                    Ok(response
105                        .data
106                        .into_iter()
107                        .zip(documents.into_iter())
108                        .map(|(embedding, document)| embeddings::Embedding {
109                            document,
110                            vec: embedding
111                                .embedding
112                                .into_iter()
113                                .filter_map(|n| n.as_f64())
114                                .collect(),
115                        })
116                        .collect())
117                }
118                ApiResponse::Error(err) => {
119                    tracing::warn!(
120                        message = %err.error,
121                        "provider returned an error response"
122                    );
123                    Err(EmbeddingError::from_http_response(
124                        status,
125                        String::from_utf8_lossy(&response_body),
126                    ))
127                }
128            }
129        } else {
130            let text = http_client::text(response).await?;
131            Err(EmbeddingError::from_http_response(status, text))
132        }
133    }
134}
135
136impl<T> EmbeddingModel<T>
137where
138    T: Default,
139{
140    pub fn new(client: Client<T>, model: impl Into<String>, ndims: usize) -> Self {
141        Self {
142            client,
143            model: model.into(),
144            ndims,
145        }
146    }
147}
rig_core/providers/together/embedding.rs

rig_core/providers/together/
embedding.rs
