Skip to content

Commit 8d9b3a9

Browse files
ojhermannojhermann
committed
refactor(providers)!: adopt GenericEmbeddingModel for together/openrouter/mistral embeddings
Collapse the hand-rolled Together, OpenRouter, and Mistral embedding models onto the shared GenericEmbeddingModel<Ext> via a new OpenAIEmbeddingsCompatible path hook, mirroring OpenAICompatibleProvider::completion_path. Together embeddings now honor dimensions (via ndims), encoding_format, and user, which the previous hand-rolled copy silently dropped. EmbeddingResponse.usage becomes optional so providers that omit it (OpenRouter for some models, Together historically) are not regressed now that they share one response type. Removes the now-dead per-provider embedding response/format types and their orphaned error envelopes. Azure (needs the path hook plus deployment-scoped URLs) and Copilot (blocked on async auth headers) are left out of scope per the issue. Fixes #2078.
1 parent f885c20 commit 8d9b3a9

9 files changed

Lines changed: 157 additions & 508 deletions

File tree

‎crates/rig-core/CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
99

1010
### Changed
1111

12+
- *(providers)* [**breaking**] Collapse the hand-rolled Together, OpenRouter, and Mistral embedding models onto the shared `GenericEmbeddingModel<Ext>`, routed through a new `OpenAIEmbeddingsCompatible` path hook (mirroring `OpenAICompatibleProvider::completion_path`). Together embeddings now honor `dimensions` (via `ndims`), `encoding_format`, and `user`, which the previous implementation silently dropped. Removes the now-redundant per-provider embedding types: `together::{EmbeddingResponse, EmbeddingData, Usage}` and the `together_ai_api_types` module, `openrouter::{EncodingFormat, EmbeddingResponse, EmbeddingData}`, and `mistral::{EmbeddingResponse, EmbeddingData}`.
1213
- *(core)* [**breaking**] Mark `PromptError`, `StructuredOutputError`, `ToolError`, `ToolSetError`, and `VectorStoreError` as non-exhaustive, requiring downstream match expressions to include a wildcard arm. Conversation memory load failures now surface as the typed `PromptError::MemoryError` variant instead of `CompletionError::RequestError`.
1314

1415
## [0.40.0](https://github.com/0xPlaygrounds/rig/compare/rig-core-v0.39.0...rig-core-v0.40.0) - 2026-07-10

‎crates/rig-core/src/providers/llamafile.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,10 @@ impl openai::completion::OpenAICompatibleProvider for LlamafileExt {
6868
type Response = openai::CompletionResponse;
6969
}
7070

71+
// llama.cpp's `build_uri` injects the `/v1/` segment, so the default
72+
// `/embeddings` path is correct.
73+
impl openai::embedding::OpenAIEmbeddingsCompatible for LlamafileExt {}
74+
7175
impl<H> Capabilities<H> for LlamafileExt {
7276
type Completion = Capable<openai::completion::GenericCompletionModel<LlamafileExt, H>>;
7377
type Embeddings = Capable<openai::embedding::GenericEmbeddingModel<LlamafileExt, H>>;

‎crates/rig-core/src/providers/mistral/client.rs‎

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,13 @@ impl crate::providers::openai::completion::OpenAICompatibleProvider for MistralE
101101
}
102102
}
103103

104+
impl crate::providers::openai::embedding::OpenAIEmbeddingsCompatible for MistralExt {
105+
// The client base URL is the bare host, so the version segment is explicit.
106+
fn embeddings_path(&self) -> String {
107+
"/v1/embeddings".to_string()
108+
}
109+
}
110+
104111
impl<H> Capabilities<H> for MistralExt {
105112
type Completion = Capable<super::CompletionModel<H>>;
106113
type Embeddings = Capable<super::EmbeddingModel<H>>;
@@ -217,18 +224,6 @@ impl std::fmt::Display for Usage {
217224
}
218225
}
219226

220-
#[derive(Debug, Deserialize)]
221-
pub struct ApiErrorResponse {
222-
pub(crate) message: String,
223-
}
224-
225-
#[derive(Debug, Deserialize)]
226-
#[serde(untagged)]
227-
pub(crate) enum ApiResponse<T> {
228-
Ok(T),
229-
Err(ApiErrorResponse),
230-
}
231-
232227
#[cfg(test)]
233228
mod tests {
234229
#[test]

‎crates/rig-core/src/providers/mistral/embedding.rs‎

Lines changed: 10 additions & 131 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
1-
use serde::Deserialize;
2-
use serde_json::json;
1+
//! Mistral's embeddings endpoint is OpenAI-compatible, so [`EmbeddingModel`] is
2+
//! a thin alias over the shared [`GenericEmbeddingModel`]. The `/v1/embeddings`
3+
//! path is supplied by [`MistralExt`](super::client::MistralExt)'s
4+
//! [`OpenAIEmbeddingsCompatible`](crate::providers::openai::embedding::OpenAIEmbeddingsCompatible)
5+
//! implementation in `client.rs`.
36
4-
use crate::{
5-
embeddings::{self, EmbeddingError},
6-
http_client::{self, HttpClientExt},
7-
};
8-
9-
use super::client::{ApiResponse, Client, Usage};
7+
use super::client::MistralExt;
8+
use crate::providers::openai::embedding::GenericEmbeddingModel;
109

1110
// ================================================================
1211
// Mistral Embedding API
@@ -15,126 +14,6 @@ pub const MISTRAL_EMBED: &str = "mistral-embed";
1514

1615
pub const MAX_DOCUMENTS: usize = 1024;
1716

18-
#[derive(Clone)]
19-
pub struct EmbeddingModel<T = reqwest::Client> {
20-
client: Client<T>,
21-
pub model: String,
22-
ndims: usize,
23-
}
24-
25-
impl<T> EmbeddingModel<T> {
26-
pub fn new(client: Client<T>, model: impl Into<String>, ndims: usize) -> Self {
27-
Self {
28-
client,
29-
model: model.into(),
30-
ndims,
31-
}
32-
}
33-
34-
pub fn with_model(client: Client<T>, model: &str, ndims: usize) -> Self {
35-
Self {
36-
client,
37-
model: model.to_string(),
38-
ndims,
39-
}
40-
}
41-
}
42-
43-
impl<T> embeddings::EmbeddingModel for EmbeddingModel<T>
44-
where
45-
T: HttpClientExt + Clone + 'static,
46-
{
47-
type Client = Client<T>;
48-
49-
const MAX_DOCUMENTS: usize = MAX_DOCUMENTS;
50-
51-
fn make(client: &Self::Client, model: impl Into<String>, dims: Option<usize>) -> Self {
52-
Self::new(client.clone(), model, dims.unwrap_or_default())
53-
}
54-
55-
fn ndims(&self) -> usize {
56-
self.ndims
57-
}
58-
59-
async fn embed_texts(
60-
&self,
61-
documents: impl IntoIterator<Item = String>,
62-
) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
63-
let documents = documents.into_iter().collect::<Vec<_>>();
64-
65-
let body = serde_json::to_vec(&json!({
66-
"model": self.model,
67-
"input": documents
68-
}))?;
69-
70-
let req = self
71-
.client
72-
.post("v1/embeddings")?
73-
.header("Content-Type", "application/json")
74-
.body(body)
75-
.map_err(|e| EmbeddingError::HttpError(e.into()))?;
76-
77-
let response = self.client.send(req).await?;
78-
79-
let status = response.status();
80-
if status.is_success() {
81-
let response_body: Vec<u8> = response.into_body().await?;
82-
let parsed: ApiResponse<EmbeddingResponse> = serde_json::from_slice(&response_body)?;
83-
84-
match parsed {
85-
ApiResponse::Ok(response) => {
86-
tracing::debug!(target: "rig",
87-
"Mistral embedding token usage: {}",
88-
response.usage
89-
);
90-
91-
if response.data.len() != documents.len() {
92-
return Err(EmbeddingError::ResponseError(
93-
"Response data length does not match input length".into(),
94-
));
95-
}
96-
97-
Ok(response
98-
.data
99-
.into_iter()
100-
.zip(documents.into_iter())
101-
.map(|(embedding, document)| embeddings::Embedding {
102-
document,
103-
vec: embedding
104-
.embedding
105-
.into_iter()
106-
.filter_map(|n| n.as_f64())
107-
.collect(),
108-
})
109-
.collect())
110-
}
111-
ApiResponse::Err(err) => {
112-
tracing::warn!(message = %err.message, "provider returned an error response");
113-
Err(EmbeddingError::from_http_response(
114-
status,
115-
String::from_utf8_lossy(&response_body),
116-
))
117-
}
118-
}
119-
} else {
120-
let text = http_client::text(response).await?;
121-
Err(EmbeddingError::from_http_response(status, text))
122-
}
123-
}
124-
}
125-
126-
#[derive(Debug, Deserialize)]
127-
pub struct EmbeddingResponse {
128-
pub id: String,
129-
pub object: String,
130-
pub model: String,
131-
pub usage: Usage,
132-
pub data: Vec<EmbeddingData>,
133-
}
134-
135-
#[derive(Debug, Deserialize)]
136-
pub struct EmbeddingData {
137-
pub object: String,
138-
pub embedding: Vec<serde_json::Number>,
139-
pub index: usize,
140-
}
17+
/// Mistral embedding model, backed by the shared OpenAI-compatible
18+
/// [`GenericEmbeddingModel`].
19+
pub type EmbeddingModel<H = reqwest::Client> = GenericEmbeddingModel<MistralExt, H>;

‎crates/rig-core/src/providers/openai/embedding.rs‎

Lines changed: 55 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,44 @@ pub struct EmbeddingResponse {
2020
pub object: String,
2121
pub data: Vec<EmbeddingData>,
2222
pub model: String,
23-
pub usage: Usage,
23+
/// Token usage for the request. Optional because some OpenAI-compatible
24+
/// providers routed through [`GenericEmbeddingModel`] omit it (e.g.
25+
/// OpenRouter returns it only for some models, Together historically did
26+
/// not surface it at all). OpenAI itself always populates it.
27+
#[serde(default)]
28+
pub usage: Option<Usage>,
2429
}
2530

31+
/// Provider hook selecting the request path for the OpenAI-compatible
32+
/// embeddings endpoint.
33+
///
34+
/// This mirrors [`OpenAICompatibleProvider::completion_path`] on the completion
35+
/// side: it lets a provider that shares the OpenAI embeddings transport but
36+
/// exposes it at a different path (typically a bare-host base URL that needs an
37+
/// explicit `/v1` segment) reuse [`GenericEmbeddingModel`] instead of
38+
/// hand-rolling the request flow.
39+
///
40+
/// It is deliberately a dedicated trait rather than a method on
41+
/// [`OpenAICompatibleProvider`](super::completion::OpenAICompatibleProvider):
42+
/// a provider's embeddings capability and chat-completions capability can live
43+
/// on different provider extensions. OpenAI's own embeddings run on
44+
/// [`OpenAIResponsesExt`](super::OpenAIResponsesExt), which is not an
45+
/// `OpenAICompatibleProvider`, so bounding the embedding model on that trait
46+
/// would exclude OpenAI itself.
47+
#[doc(hidden)]
48+
pub trait OpenAIEmbeddingsCompatible: crate::client::Provider {
49+
/// The request path for embeddings, resolved against the client base URL by
50+
/// [`Provider::build_uri`](crate::client::Provider::build_uri). Defaults to
51+
/// `/embeddings`; providers whose base URL omits the API-version segment
52+
/// override this (e.g. `/v1/embeddings`).
53+
fn embeddings_path(&self) -> String {
54+
"/embeddings".to_string()
55+
}
56+
}
57+
58+
impl OpenAIEmbeddingsCompatible for super::OpenAIResponsesExt {}
59+
impl OpenAIEmbeddingsCompatible for super::OpenAICompletionsExt {}
60+
2661
#[derive(Debug, Deserialize, Clone, Serialize)]
2762
#[serde(rename_all = "snake_case")]
2863
pub enum EncodingFormat {
@@ -64,7 +99,7 @@ fn model_dimensions_from_identifier(identifier: &str) -> Option<usize> {
6499
impl<Ext, H> embeddings::EmbeddingModel for GenericEmbeddingModel<Ext, H>
65100
where
66101
crate::client::Client<Ext, H>: HttpClientExt + Clone + std::fmt::Debug + Send + 'static,
67-
Ext: crate::client::Provider + Clone + 'static,
102+
Ext: OpenAIEmbeddingsCompatible + Clone + 'static,
68103
H: Clone + Default + std::fmt::Debug + 'static,
69104
{
70105
const MAX_DOCUMENTS: usize = 1024;
@@ -124,7 +159,7 @@ where
124159

125160
let req = self
126161
.client
127-
.post("/embeddings")?
162+
.post(self.client.ext().embeddings_path())?
128163
.body(body)
129164
.map_err(|e| EmbeddingError::HttpError(e.into()))?;
130165

@@ -138,7 +173,7 @@ where
138173
match parsed {
139174
ApiResponse::Ok(response) => {
140175
tracing::info!(target: "rig",
141-
"OpenAI embedding token usage: {:?}",
176+
"embedding token usage: {:?}",
142177
response.usage
143178
);
144179

@@ -148,19 +183,22 @@ where
148183
));
149184
}
150185

151-
let usage = crate::completion::Usage {
152-
input_tokens: response.usage.prompt_tokens as u64,
153-
output_tokens: 0,
154-
total_tokens: response.usage.total_tokens as u64,
155-
cached_input_tokens: response
156-
.usage
157-
.prompt_tokens_details
158-
.as_ref()
159-
.map_or(0, |d| d.cached_tokens as u64),
160-
cache_creation_input_tokens: 0,
161-
tool_use_prompt_tokens: 0,
162-
reasoning_tokens: 0,
163-
};
186+
let usage = response
187+
.usage
188+
.as_ref()
189+
.map(|usage| crate::completion::Usage {
190+
input_tokens: usage.prompt_tokens as u64,
191+
output_tokens: 0,
192+
total_tokens: usage.total_tokens as u64,
193+
cached_input_tokens: usage
194+
.prompt_tokens_details
195+
.as_ref()
196+
.map_or(0, |d| d.cached_tokens as u64),
197+
cache_creation_input_tokens: 0,
198+
tool_use_prompt_tokens: 0,
199+
reasoning_tokens: 0,
200+
})
201+
.unwrap_or_else(crate::completion::Usage::new);
164202

165203
let embeddings = response
166204
.data

‎crates/rig-core/src/providers/openrouter/client.rs‎

Lines changed: 4 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,10 @@ impl<H> Capabilities<H> for OpenRouterExt {
4545
type Rerank = Nothing;
4646
}
4747

48+
// OpenRouter's base URL already includes `/api/v1`, so the default
49+
// `/embeddings` path is correct.
50+
impl crate::providers::openai::embedding::OpenAIEmbeddingsCompatible for OpenRouterExt {}
51+
4852
impl DebugExt for OpenRouterExt {}
4953

5054
impl ProviderBuilder for OpenRouterExtBuilder {
@@ -82,18 +86,6 @@ impl ProviderClient for Client {
8286
}
8387
}
8488

85-
#[derive(Debug, Deserialize)]
86-
pub(crate) struct ApiErrorResponse {
87-
pub message: String,
88-
}
89-
90-
#[derive(Debug, Deserialize)]
91-
#[serde(untagged)]
92-
pub(crate) enum ApiResponse<T> {
93-
Ok(T),
94-
Err(ApiErrorResponse),
95-
}
96-
9789
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
9890
pub struct Usage {
9991
pub prompt_tokens: usize,

0 commit comments

Comments
 (0)