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

rig/providers/together/
embedding.rs
