From 8d9b3a935d4d60ab3d3402eab34eaf751563a0cf Mon Sep 17 00:00:00 2001 From: Otto Hermann Date: Sun, 12 Jul 2026 15:56:32 +0000 Subject: [PATCH 1/2] refactor(providers)!: adopt GenericEmbeddingModel for together/openrouter/mistral embeddings Collapse the hand-rolled Together, OpenRouter, and Mistral embedding models onto the shared GenericEmbeddingModel 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. --- crates/rig-core/CHANGELOG.md | 1 + crates/rig-core/src/providers/llamafile.rs | 4 + .../rig-core/src/providers/mistral/client.rs | 19 +- .../src/providers/mistral/embedding.rs | 141 +----------- .../src/providers/openai/embedding.rs | 72 +++++-- .../src/providers/openrouter/client.rs | 16 +- .../src/providers/openrouter/embedding.rs | 202 ++---------------- .../rig-core/src/providers/together/client.rs | 29 +-- .../src/providers/together/embedding.rs | 181 +++++----------- 9 files changed, 157 insertions(+), 508 deletions(-) diff --git a/crates/rig-core/CHANGELOG.md b/crates/rig-core/CHANGELOG.md index aff615c91..f774cfdcd 100644 --- a/crates/rig-core/CHANGELOG.md +++ b/crates/rig-core/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- *(providers)* [**breaking**] Collapse the hand-rolled Together, OpenRouter, and Mistral embedding models onto the shared `GenericEmbeddingModel`, 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}`. - *(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`. ## [0.40.0](https://github.com/0xPlaygrounds/rig/compare/rig-core-v0.39.0...rig-core-v0.40.0) - 2026-07-10 diff --git a/crates/rig-core/src/providers/llamafile.rs b/crates/rig-core/src/providers/llamafile.rs index 4fe97ee0e..4b83210cb 100644 --- a/crates/rig-core/src/providers/llamafile.rs +++ b/crates/rig-core/src/providers/llamafile.rs @@ -68,6 +68,10 @@ impl openai::completion::OpenAICompatibleProvider for LlamafileExt { type Response = openai::CompletionResponse; } +// llama.cpp's `build_uri` injects the `/v1/` segment, so the default +// `/embeddings` path is correct. +impl openai::embedding::OpenAIEmbeddingsCompatible for LlamafileExt {} + impl Capabilities for LlamafileExt { type Completion = Capable>; type Embeddings = Capable>; diff --git a/crates/rig-core/src/providers/mistral/client.rs b/crates/rig-core/src/providers/mistral/client.rs index 22f33dcba..a3acb68ad 100644 --- a/crates/rig-core/src/providers/mistral/client.rs +++ b/crates/rig-core/src/providers/mistral/client.rs @@ -101,6 +101,13 @@ impl crate::providers::openai::completion::OpenAICompatibleProvider for MistralE } } +impl crate::providers::openai::embedding::OpenAIEmbeddingsCompatible for MistralExt { + // The client base URL is the bare host, so the version segment is explicit. + fn embeddings_path(&self) -> String { + "/v1/embeddings".to_string() + } +} + impl Capabilities for MistralExt { type Completion = Capable>; type Embeddings = Capable>; @@ -217,18 +224,6 @@ impl std::fmt::Display for Usage { } } -#[derive(Debug, Deserialize)] -pub struct ApiErrorResponse { - pub(crate) message: String, -} - -#[derive(Debug, Deserialize)] -#[serde(untagged)] -pub(crate) enum ApiResponse { - Ok(T), - Err(ApiErrorResponse), -} - #[cfg(test)] mod tests { #[test] diff --git a/crates/rig-core/src/providers/mistral/embedding.rs b/crates/rig-core/src/providers/mistral/embedding.rs index c6b366ded..8997e769a 100644 --- a/crates/rig-core/src/providers/mistral/embedding.rs +++ b/crates/rig-core/src/providers/mistral/embedding.rs @@ -1,12 +1,11 @@ -use serde::Deserialize; -use serde_json::json; +//! Mistral's embeddings endpoint is OpenAI-compatible, so [`EmbeddingModel`] is +//! a thin alias over the shared [`GenericEmbeddingModel`]. The `/v1/embeddings` +//! path is supplied by [`MistralExt`](super::client::MistralExt)'s +//! [`OpenAIEmbeddingsCompatible`](crate::providers::openai::embedding::OpenAIEmbeddingsCompatible) +//! implementation in `client.rs`. -use crate::{ - embeddings::{self, EmbeddingError}, - http_client::{self, HttpClientExt}, -}; - -use super::client::{ApiResponse, Client, Usage}; +use super::client::MistralExt; +use crate::providers::openai::embedding::GenericEmbeddingModel; // ================================================================ // Mistral Embedding API @@ -15,126 +14,6 @@ pub const MISTRAL_EMBED: &str = "mistral-embed"; pub const MAX_DOCUMENTS: usize = 1024; -#[derive(Clone)] -pub struct EmbeddingModel { - client: Client, - pub model: String, - ndims: usize, -} - -impl EmbeddingModel { - pub fn new(client: Client, model: impl Into, ndims: usize) -> Self { - Self { - client, - model: model.into(), - ndims, - } - } - - pub fn with_model(client: Client, model: &str, ndims: usize) -> Self { - Self { - client, - model: model.to_string(), - ndims, - } - } -} - -impl embeddings::EmbeddingModel for EmbeddingModel -where - T: HttpClientExt + Clone + 'static, -{ - type Client = Client; - - const MAX_DOCUMENTS: usize = MAX_DOCUMENTS; - - fn make(client: &Self::Client, model: impl Into, dims: Option) -> Self { - Self::new(client.clone(), model, dims.unwrap_or_default()) - } - - fn ndims(&self) -> usize { - self.ndims - } - - async fn embed_texts( - &self, - documents: impl IntoIterator, - ) -> Result, EmbeddingError> { - let documents = documents.into_iter().collect::>(); - - let body = serde_json::to_vec(&json!({ - "model": self.model, - "input": documents - }))?; - - let req = self - .client - .post("v1/embeddings")? - .header("Content-Type", "application/json") - .body(body) - .map_err(|e| EmbeddingError::HttpError(e.into()))?; - - let response = self.client.send(req).await?; - - let status = response.status(); - if status.is_success() { - let response_body: Vec = response.into_body().await?; - let parsed: ApiResponse = serde_json::from_slice(&response_body)?; - - match parsed { - ApiResponse::Ok(response) => { - tracing::debug!(target: "rig", - "Mistral embedding token usage: {}", - response.usage - ); - - if response.data.len() != documents.len() { - return Err(EmbeddingError::ResponseError( - "Response data length does not match input length".into(), - )); - } - - Ok(response - .data - .into_iter() - .zip(documents.into_iter()) - .map(|(embedding, document)| embeddings::Embedding { - document, - vec: embedding - .embedding - .into_iter() - .filter_map(|n| n.as_f64()) - .collect(), - }) - .collect()) - } - ApiResponse::Err(err) => { - tracing::warn!(message = %err.message, "provider returned an error response"); - Err(EmbeddingError::from_http_response( - status, - String::from_utf8_lossy(&response_body), - )) - } - } - } else { - let text = http_client::text(response).await?; - Err(EmbeddingError::from_http_response(status, text)) - } - } -} - -#[derive(Debug, Deserialize)] -pub struct EmbeddingResponse { - pub id: String, - pub object: String, - pub model: String, - pub usage: Usage, - pub data: Vec, -} - -#[derive(Debug, Deserialize)] -pub struct EmbeddingData { - pub object: String, - pub embedding: Vec, - pub index: usize, -} +/// Mistral embedding model, backed by the shared OpenAI-compatible +/// [`GenericEmbeddingModel`]. +pub type EmbeddingModel = GenericEmbeddingModel; diff --git a/crates/rig-core/src/providers/openai/embedding.rs b/crates/rig-core/src/providers/openai/embedding.rs index 058b7643f..0c46fbd09 100644 --- a/crates/rig-core/src/providers/openai/embedding.rs +++ b/crates/rig-core/src/providers/openai/embedding.rs @@ -20,9 +20,44 @@ pub struct EmbeddingResponse { pub object: String, pub data: Vec, pub model: String, - pub usage: Usage, + /// Token usage for the request. Optional because some OpenAI-compatible + /// providers routed through [`GenericEmbeddingModel`] omit it (e.g. + /// OpenRouter returns it only for some models, Together historically did + /// not surface it at all). OpenAI itself always populates it. + #[serde(default)] + pub usage: Option, } +/// Provider hook selecting the request path for the OpenAI-compatible +/// embeddings endpoint. +/// +/// This mirrors [`OpenAICompatibleProvider::completion_path`] on the completion +/// side: it lets a provider that shares the OpenAI embeddings transport but +/// exposes it at a different path (typically a bare-host base URL that needs an +/// explicit `/v1` segment) reuse [`GenericEmbeddingModel`] instead of +/// hand-rolling the request flow. +/// +/// It is deliberately a dedicated trait rather than a method on +/// [`OpenAICompatibleProvider`](super::completion::OpenAICompatibleProvider): +/// a provider's embeddings capability and chat-completions capability can live +/// on different provider extensions. OpenAI's own embeddings run on +/// [`OpenAIResponsesExt`](super::OpenAIResponsesExt), which is not an +/// `OpenAICompatibleProvider`, so bounding the embedding model on that trait +/// would exclude OpenAI itself. +#[doc(hidden)] +pub trait OpenAIEmbeddingsCompatible: crate::client::Provider { + /// The request path for embeddings, resolved against the client base URL by + /// [`Provider::build_uri`](crate::client::Provider::build_uri). Defaults to + /// `/embeddings`; providers whose base URL omits the API-version segment + /// override this (e.g. `/v1/embeddings`). + fn embeddings_path(&self) -> String { + "/embeddings".to_string() + } +} + +impl OpenAIEmbeddingsCompatible for super::OpenAIResponsesExt {} +impl OpenAIEmbeddingsCompatible for super::OpenAICompletionsExt {} + #[derive(Debug, Deserialize, Clone, Serialize)] #[serde(rename_all = "snake_case")] pub enum EncodingFormat { @@ -64,7 +99,7 @@ fn model_dimensions_from_identifier(identifier: &str) -> Option { impl embeddings::EmbeddingModel for GenericEmbeddingModel where crate::client::Client: HttpClientExt + Clone + std::fmt::Debug + Send + 'static, - Ext: crate::client::Provider + Clone + 'static, + Ext: OpenAIEmbeddingsCompatible + Clone + 'static, H: Clone + Default + std::fmt::Debug + 'static, { const MAX_DOCUMENTS: usize = 1024; @@ -124,7 +159,7 @@ where let req = self .client - .post("/embeddings")? + .post(self.client.ext().embeddings_path())? .body(body) .map_err(|e| EmbeddingError::HttpError(e.into()))?; @@ -138,7 +173,7 @@ where match parsed { ApiResponse::Ok(response) => { tracing::info!(target: "rig", - "OpenAI embedding token usage: {:?}", + "embedding token usage: {:?}", response.usage ); @@ -148,19 +183,22 @@ where )); } - let usage = crate::completion::Usage { - input_tokens: response.usage.prompt_tokens as u64, - output_tokens: 0, - total_tokens: response.usage.total_tokens as u64, - cached_input_tokens: response - .usage - .prompt_tokens_details - .as_ref() - .map_or(0, |d| d.cached_tokens as u64), - cache_creation_input_tokens: 0, - tool_use_prompt_tokens: 0, - reasoning_tokens: 0, - }; + let usage = response + .usage + .as_ref() + .map(|usage| crate::completion::Usage { + input_tokens: usage.prompt_tokens as u64, + output_tokens: 0, + total_tokens: usage.total_tokens as u64, + cached_input_tokens: usage + .prompt_tokens_details + .as_ref() + .map_or(0, |d| d.cached_tokens as u64), + cache_creation_input_tokens: 0, + tool_use_prompt_tokens: 0, + reasoning_tokens: 0, + }) + .unwrap_or_else(crate::completion::Usage::new); let embeddings = response .data diff --git a/crates/rig-core/src/providers/openrouter/client.rs b/crates/rig-core/src/providers/openrouter/client.rs index fd22fa795..a81bd5257 100644 --- a/crates/rig-core/src/providers/openrouter/client.rs +++ b/crates/rig-core/src/providers/openrouter/client.rs @@ -45,6 +45,10 @@ impl Capabilities for OpenRouterExt { type Rerank = Nothing; } +// OpenRouter's base URL already includes `/api/v1`, so the default +// `/embeddings` path is correct. +impl crate::providers::openai::embedding::OpenAIEmbeddingsCompatible for OpenRouterExt {} + impl DebugExt for OpenRouterExt {} impl ProviderBuilder for OpenRouterExtBuilder { @@ -82,18 +86,6 @@ impl ProviderClient for Client { } } -#[derive(Debug, Deserialize)] -pub(crate) struct ApiErrorResponse { - pub message: String, -} - -#[derive(Debug, Deserialize)] -#[serde(untagged)] -pub(crate) enum ApiResponse { - Ok(T), - Err(ApiErrorResponse), -} - #[derive(Clone, Debug, Default, Deserialize, Serialize)] pub struct Usage { pub prompt_tokens: usize, diff --git a/crates/rig-core/src/providers/openrouter/embedding.rs b/crates/rig-core/src/providers/openrouter/embedding.rs index 0e86e13cc..bff1ca71b 100644 --- a/crates/rig-core/src/providers/openrouter/embedding.rs +++ b/crates/rig-core/src/providers/openrouter/embedding.rs @@ -1,189 +1,13 @@ -use super::{Client, Usage, client::ApiResponse}; -use crate::embeddings::EmbeddingError; -use crate::http_client::HttpClientExt; -use crate::wasm_compat::WasmCompatSend; -use crate::{embeddings, http_client}; -use serde::{Deserialize, Serialize}; -use serde_json::json; - -#[derive(Debug, Deserialize)] -pub struct EmbeddingResponse { - pub object: String, - pub data: Vec, - pub model: String, - pub usage: Option, - pub id: Option, -} - -#[derive(Debug, Deserialize, Clone, Serialize)] -#[serde(rename_all = "snake_case")] -pub enum EncodingFormat { - Float, - Base64, -} - -#[derive(Debug, Deserialize)] -pub struct EmbeddingData { - pub object: String, - pub embedding: Vec, - pub index: usize, -} - -#[derive(Clone)] -pub struct EmbeddingModel { - client: Client, - pub model: String, - pub encoding_format: Option, - pub user: Option, - ndims: usize, -} - -impl embeddings::EmbeddingModel for EmbeddingModel -where - T: HttpClientExt + Clone + std::fmt::Debug + Default + WasmCompatSend + 'static, -{ - const MAX_DOCUMENTS: usize = 1024; - - type Client = Client; - - fn make(client: &Self::Client, model: impl Into, ndims: Option) -> Self { - let model = model.into(); - let dims = ndims.unwrap_or_default(); - - Self::new(client.clone(), model, dims) - } - - fn ndims(&self) -> usize { - self.ndims - } - - async fn embed_texts( - &self, - documents: impl IntoIterator, - ) -> Result, EmbeddingError> { - let documents = documents.into_iter().collect::>(); - - let mut body = json!({ - "model": self.model, - "input": documents, - }); - - let body_object = body.as_object_mut().ok_or_else(|| { - EmbeddingError::ResponseError("embedding request body must be a JSON object".into()) - })?; - - if self.ndims > 0 { - body_object.insert("dimensions".to_owned(), json!(self.ndims)); - } - - if let Some(encoding_format) = &self.encoding_format { - body_object.insert("encoding_format".to_owned(), json!(encoding_format)); - } - - if let Some(user) = &self.user { - body_object.insert("user".to_owned(), json!(user)); - } - - let body = serde_json::to_vec(&body)?; - - let req = self - .client - .post("/embeddings")? - .body(body) - .map_err(|e| EmbeddingError::HttpError(e.into()))?; - - let response = self.client.send(req).await?; - - let status = response.status(); - if status.is_success() { - let response_body: Vec = response.into_body().await?; - let parsed: ApiResponse = serde_json::from_slice(&response_body)?; - - match parsed { - ApiResponse::Ok(response) => { - tracing::info!(target: "rig", - "OpenRouter embedding token usage: {:?}", - response.usage - ); - - if response.data.len() != documents.len() { - return Err(EmbeddingError::ResponseError( - "Response data length does not match input length".into(), - )); - } - - Ok(response - .data - .into_iter() - .zip(documents.into_iter()) - .map(|(embedding, document)| embeddings::Embedding { - document, - vec: embedding - .embedding - .into_iter() - .filter_map(|n| n.as_f64()) - .collect(), - }) - .collect()) - } - ApiResponse::Err(err) => { - tracing::warn!(message = %err.message, "provider returned an error response"); - Err(EmbeddingError::from_http_response( - status, - String::from_utf8_lossy(&response_body).into_owned(), - )) - } - } - } else { - let text = http_client::text(response).await?; - Err(EmbeddingError::from_http_response(status, text)) - } - } -} - -impl EmbeddingModel { - pub fn new(client: Client, model: impl Into, ndims: usize) -> Self { - Self { - client, - model: model.into(), - encoding_format: None, - ndims, - user: None, - } - } - - pub fn with_model(client: Client, model: &str, ndims: usize) -> Self { - Self { - client, - model: model.into(), - encoding_format: None, - ndims, - user: None, - } - } - - pub fn with_encoding_format( - client: Client, - model: &str, - ndims: usize, - encoding_format: EncodingFormat, - ) -> Self { - Self { - client, - model: model.into(), - encoding_format: Some(encoding_format), - ndims, - user: None, - } - } - - pub fn encoding_format(mut self, encoding_format: EncodingFormat) -> Self { - self.encoding_format = Some(encoding_format); - self - } - - pub fn user(mut self, user: impl Into) -> Self { - self.user = Some(user.into()); - self - } -} +//! OpenRouter's embeddings endpoint is OpenAI-compatible, so [`EmbeddingModel`] +//! is a thin alias over the shared [`GenericEmbeddingModel`]. OpenRouter's base +//! URL already carries the `/api/v1` segment, so the default `/embeddings` path +//! from the +//! [`OpenAIEmbeddingsCompatible`](crate::providers::openai::embedding::OpenAIEmbeddingsCompatible) +//! implementation on [`OpenRouterExt`](super::client::OpenRouterExt) is used. + +use super::client::OpenRouterExt; +use crate::providers::openai::embedding::GenericEmbeddingModel; + +/// OpenRouter embedding model, backed by the shared OpenAI-compatible +/// [`GenericEmbeddingModel`]. +pub type EmbeddingModel = GenericEmbeddingModel; diff --git a/crates/rig-core/src/providers/together/client.rs b/crates/rig-core/src/providers/together/client.rs index aa1ccf6d7..f73f00085 100644 --- a/crates/rig-core/src/providers/together/client.rs +++ b/crates/rig-core/src/providers/together/client.rs @@ -47,6 +47,13 @@ impl crate::providers::openai::completion::OpenAICompatibleProvider for Together } } +impl crate::providers::openai::embedding::OpenAIEmbeddingsCompatible for TogetherExt { + // The client base URL is the bare host, so the version segment is explicit. + fn embeddings_path(&self) -> String { + "/v1/embeddings".to_string() + } +} + impl Capabilities for TogetherExt { type Completion = Capable>; type Embeddings = Capable>; @@ -94,28 +101,6 @@ impl ProviderClient for Client { } } -pub mod together_ai_api_types { - use serde::Deserialize; - - impl ApiErrorResponse { - pub fn message(&self) -> String { - format!("Code `{}`: {}", self.code, self.error) - } - } - - #[derive(Debug, Deserialize)] - pub struct ApiErrorResponse { - pub error: String, - pub code: String, - } - - #[derive(Debug, Deserialize)] - #[serde(untagged)] - pub enum ApiResponse { - Ok(T), - Error(ApiErrorResponse), - } -} #[cfg(test)] mod tests { #[test] diff --git a/crates/rig-core/src/providers/together/embedding.rs b/crates/rig-core/src/providers/together/embedding.rs index c59457f58..837228f71 100644 --- a/crates/rig-core/src/providers/together/embedding.rs +++ b/crates/rig-core/src/providers/together/embedding.rs @@ -3,15 +3,14 @@ //! From [Together AI Reference](https://docs.together.ai/docs/embeddings-overview) // ================================================================ -use serde::Deserialize; -use serde_json::json; +//! Together's embeddings endpoint is OpenAI-compatible, so [`EmbeddingModel`] is +//! a thin alias over the shared [`GenericEmbeddingModel`]. The `/v1/embeddings` +//! path is supplied by [`TogetherExt`](super::client::TogetherExt)'s +//! [`OpenAIEmbeddingsCompatible`](crate::providers::openai::embedding::OpenAIEmbeddingsCompatible) +//! implementation in `client.rs`. -use crate::{ - embeddings::{self, EmbeddingError}, - http_client::{self, HttpClientExt}, -}; - -use super::{Client, client::together_ai_api_types::ApiResponse}; +use super::client::TogetherExt; +use crate::providers::openai::embedding::GenericEmbeddingModel; // ================================================================ // Together AI Embedding API @@ -26,122 +25,54 @@ pub const M2_BERT_80M_8K_RETRIEVAL: &str = "togethercomputer/m2-bert-80M-8k-retr pub const SENTENCE_BERT: &str = "sentence-transformers/msmarco-bert-base-dot-v5"; pub const UAE_LARGE_V1: &str = "WhereIsAI/UAE-Large-V1"; -#[derive(Debug, Deserialize)] -pub struct EmbeddingResponse { - pub model: String, - pub object: String, - pub data: Vec, -} - -#[derive(Debug, Deserialize)] -pub struct EmbeddingData { - pub object: String, - pub embedding: Vec, - pub index: usize, -} - -#[derive(Debug, Deserialize)] -pub struct Usage { - pub prompt_tokens: usize, - pub total_tokens: usize, -} - -#[derive(Clone)] -pub struct EmbeddingModel { - client: Client, - pub model: String, - ndims: usize, -} - -impl embeddings::EmbeddingModel for EmbeddingModel -where - T: HttpClientExt + Default + Clone + Send + 'static, -{ - const MAX_DOCUMENTS: usize = 1024; // This might need to be adjusted based on Together AI's actual limit - - type Client = Client; - - fn make(client: &Self::Client, model: impl Into, dims: Option) -> Self { - Self::new(client.clone(), model, dims.unwrap_or_default()) - } - - fn ndims(&self) -> usize { - self.ndims - } - - async fn embed_texts( - &self, - documents: impl IntoIterator, - ) -> Result, EmbeddingError> { - let documents = documents.into_iter().collect::>(); - - let body = serde_json::to_vec(&json!({ - "model": self.model, - "input": documents, - }))?; - - let req = self - .client - .post("/v1/embeddings")? - .body(body) - .map_err(|e| EmbeddingError::HttpError(e.into()))?; - - let response = self.client.send(req).await?; - - let status = response.status(); - if status.is_success() { - let response_body: Vec = response.into_body().await?; - let parsed: ApiResponse = serde_json::from_slice(&response_body)?; - - match parsed { - ApiResponse::Ok(response) => { - if response.data.len() != documents.len() { - return Err(EmbeddingError::ResponseError( - "Response data length does not match input length".into(), - )); - } - - Ok(response - .data - .into_iter() - .zip(documents.into_iter()) - .map(|(embedding, document)| embeddings::Embedding { - document, - vec: embedding - .embedding - .into_iter() - .filter_map(|n| n.as_f64()) - .collect(), - }) - .collect()) - } - ApiResponse::Error(err) => { - tracing::warn!( - message = %err.error, - "provider returned an error response" - ); - Err(EmbeddingError::from_http_response( - status, - String::from_utf8_lossy(&response_body), - )) - } - } - } else { - let text = http_client::text(response).await?; - Err(EmbeddingError::from_http_response(status, text)) - } - } -} - -impl EmbeddingModel -where - T: Default, -{ - pub fn new(client: Client, model: impl Into, ndims: usize) -> Self { - Self { - client, - model: model.into(), - ndims, - } +/// Together AI embedding model, backed by the shared OpenAI-compatible +/// [`GenericEmbeddingModel`]. Unlike the previous hand-rolled implementation +/// this honors `ndims` (sent as `dimensions`), `encoding_format`, and `user`. +pub type EmbeddingModel = GenericEmbeddingModel; + +#[cfg(test)] +mod tests { + use super::BGE_BASE_EN_V1_5; + use crate::client::EmbeddingsClient; + use crate::embeddings::EmbeddingModel as _; + use crate::providers::together; + use crate::test_utils::RecordingHttpClient; + + // Minimal OpenAI-compatible embeddings response. `usage` is intentionally + // omitted to exercise the optional-usage path the generic model now allows. + const RESPONSE_BODY: &str = r#"{ + "object": "list", + "model": "BAAI/bge-base-en-v1.5", + "data": [{ "object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3] }] + }"#; + + #[tokio::test] + async fn together_embeddings_send_dimensions_to_v1_path() { + let http_client = RecordingHttpClient::new(RESPONSE_BODY); + let client = together::Client::builder() + .api_key("dummy-key") + .http_client(http_client.clone()) + .build() + .expect("client should build"); + + // With `ndims` set, `dimensions` must be sent — the previous hand-rolled + // Together model silently dropped it (the bug this migration fixes). + let model = client.embedding_model_with_ndims(BGE_BASE_EN_V1_5, 3); + model + .embed_texts(["hello".to_string()]) + .await + .expect("embedding request should succeed"); + + let requests = http_client.requests(); + assert_eq!(requests.len(), 1); + assert!( + requests[0].uri.ends_with("/v1/embeddings"), + "expected Together to POST /v1/embeddings, got {}", + requests[0].uri + ); + let body: serde_json::Value = + serde_json::from_slice(&requests[0].body).expect("request body should be JSON"); + assert_eq!(body["dimensions"], serde_json::json!(3)); + assert_eq!(body["model"], BGE_BASE_EN_V1_5); } } From f71e96a4242d5f8a0c55f426a46df99dfa83b677 Mon Sep 17 00:00:00 2001 From: Otto Hermann Date: Sun, 12 Jul 2026 16:15:00 +0000 Subject: [PATCH 2/2] test(providers): add embedding path/usage coverage; trim docs to repo conventions Add unit tests covering per-provider embeddings_path resolution and the optional-usage handling: - mistral: asserts POST /v1/embeddings, embedding parse, and usage mapping - openrouter: asserts POST .../api/v1/embeddings and zeroed usage when absent - together: adds a negative test that `dimensions` is omitted when ndims unset Also trim the module/alias docs and inline comments added on the migrated files to match each provider's existing convention (mirroring the completion.rs siblings), per review feedback. --- crates/rig-core/src/providers/llamafile.rs | 2 - .../rig-core/src/providers/mistral/client.rs | 1 - .../src/providers/mistral/embedding.rs | 57 +++++++++++++++---- .../src/providers/openai/embedding.rs | 24 +------- .../src/providers/openrouter/client.rs | 2 - .../src/providers/openrouter/embedding.rs | 47 ++++++++++++--- .../rig-core/src/providers/together/client.rs | 1 - .../src/providers/together/embedding.rs | 44 ++++++++------ 8 files changed, 113 insertions(+), 65 deletions(-) diff --git a/crates/rig-core/src/providers/llamafile.rs b/crates/rig-core/src/providers/llamafile.rs index 4b83210cb..7d3d4ffa5 100644 --- a/crates/rig-core/src/providers/llamafile.rs +++ b/crates/rig-core/src/providers/llamafile.rs @@ -68,8 +68,6 @@ impl openai::completion::OpenAICompatibleProvider for LlamafileExt { type Response = openai::CompletionResponse; } -// llama.cpp's `build_uri` injects the `/v1/` segment, so the default -// `/embeddings` path is correct. impl openai::embedding::OpenAIEmbeddingsCompatible for LlamafileExt {} impl Capabilities for LlamafileExt { diff --git a/crates/rig-core/src/providers/mistral/client.rs b/crates/rig-core/src/providers/mistral/client.rs index a3acb68ad..21c2a9be8 100644 --- a/crates/rig-core/src/providers/mistral/client.rs +++ b/crates/rig-core/src/providers/mistral/client.rs @@ -102,7 +102,6 @@ impl crate::providers::openai::completion::OpenAICompatibleProvider for MistralE } impl crate::providers::openai::embedding::OpenAIEmbeddingsCompatible for MistralExt { - // The client base URL is the bare host, so the version segment is explicit. fn embeddings_path(&self) -> String { "/v1/embeddings".to_string() } diff --git a/crates/rig-core/src/providers/mistral/embedding.rs b/crates/rig-core/src/providers/mistral/embedding.rs index 8997e769a..d93649e03 100644 --- a/crates/rig-core/src/providers/mistral/embedding.rs +++ b/crates/rig-core/src/providers/mistral/embedding.rs @@ -1,19 +1,54 @@ -//! Mistral's embeddings endpoint is OpenAI-compatible, so [`EmbeddingModel`] is -//! a thin alias over the shared [`GenericEmbeddingModel`]. The `/v1/embeddings` -//! path is supplied by [`MistralExt`](super::client::MistralExt)'s -//! [`OpenAIEmbeddingsCompatible`](crate::providers::openai::embedding::OpenAIEmbeddingsCompatible) -//! implementation in `client.rs`. - use super::client::MistralExt; use crate::providers::openai::embedding::GenericEmbeddingModel; -// ================================================================ -// Mistral Embedding API -// ================================================================ pub const MISTRAL_EMBED: &str = "mistral-embed"; pub const MAX_DOCUMENTS: usize = 1024; -/// Mistral embedding model, backed by the shared OpenAI-compatible -/// [`GenericEmbeddingModel`]. pub type EmbeddingModel = GenericEmbeddingModel; + +#[cfg(test)] +mod tests { + use super::MISTRAL_EMBED; + use crate::client::EmbeddingsClient; + use crate::embeddings::EmbeddingModel as _; + use crate::providers::mistral; + use crate::test_utils::RecordingHttpClient; + + const RESPONSE_BODY: &str = r#"{ + "id": "emb-1", + "object": "list", + "model": "mistral-embed", + "usage": { "prompt_tokens": 5, "total_tokens": 5 }, + "data": [{ "object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3] }] + }"#; + + #[tokio::test] + async fn mistral_embeddings_use_v1_path_and_map_usage() { + let http_client = RecordingHttpClient::new(RESPONSE_BODY); + let client = mistral::Client::builder() + .api_key("dummy-key") + .http_client(http_client.clone()) + .build() + .expect("client should build"); + + let model = client.embedding_model(MISTRAL_EMBED); + let response = model + .embed_texts_with_usage(["hello".to_string()]) + .await + .expect("embedding request should succeed"); + + assert_eq!(response.embeddings.len(), 1); + assert_eq!(response.embeddings[0].vec, vec![0.1, 0.2, 0.3]); + assert_eq!(response.usage.input_tokens, 5); + assert_eq!(response.usage.total_tokens, 5); + + let requests = http_client.requests(); + assert_eq!(requests.len(), 1); + assert!( + requests[0].uri.ends_with("/v1/embeddings"), + "expected Mistral to POST /v1/embeddings, got {}", + requests[0].uri + ); + } +} diff --git a/crates/rig-core/src/providers/openai/embedding.rs b/crates/rig-core/src/providers/openai/embedding.rs index 0c46fbd09..a0e8a4783 100644 --- a/crates/rig-core/src/providers/openai/embedding.rs +++ b/crates/rig-core/src/providers/openai/embedding.rs @@ -20,35 +20,17 @@ pub struct EmbeddingResponse { pub object: String, pub data: Vec, pub model: String, - /// Token usage for the request. Optional because some OpenAI-compatible - /// providers routed through [`GenericEmbeddingModel`] omit it (e.g. - /// OpenRouter returns it only for some models, Together historically did - /// not surface it at all). OpenAI itself always populates it. #[serde(default)] pub usage: Option, } -/// Provider hook selecting the request path for the OpenAI-compatible -/// embeddings endpoint. -/// -/// This mirrors [`OpenAICompatibleProvider::completion_path`] on the completion -/// side: it lets a provider that shares the OpenAI embeddings transport but -/// exposes it at a different path (typically a bare-host base URL that needs an -/// explicit `/v1` segment) reuse [`GenericEmbeddingModel`] instead of -/// hand-rolling the request flow. -/// -/// It is deliberately a dedicated trait rather than a method on -/// [`OpenAICompatibleProvider`](super::completion::OpenAICompatibleProvider): -/// a provider's embeddings capability and chat-completions capability can live -/// on different provider extensions. OpenAI's own embeddings run on -/// [`OpenAIResponsesExt`](super::OpenAIResponsesExt), which is not an -/// `OpenAICompatibleProvider`, so bounding the embedding model on that trait -/// would exclude OpenAI itself. +/// Provider hook selecting the request path for the OpenAI-compatible embeddings +/// endpoint, mirroring the completion side's `completion_path` hook. #[doc(hidden)] pub trait OpenAIEmbeddingsCompatible: crate::client::Provider { /// The request path for embeddings, resolved against the client base URL by /// [`Provider::build_uri`](crate::client::Provider::build_uri). Defaults to - /// `/embeddings`; providers whose base URL omits the API-version segment + /// `/embeddings`; providers whose base URL omits the version segment /// override this (e.g. `/v1/embeddings`). fn embeddings_path(&self) -> String { "/embeddings".to_string() diff --git a/crates/rig-core/src/providers/openrouter/client.rs b/crates/rig-core/src/providers/openrouter/client.rs index a81bd5257..8d9d2b55e 100644 --- a/crates/rig-core/src/providers/openrouter/client.rs +++ b/crates/rig-core/src/providers/openrouter/client.rs @@ -45,8 +45,6 @@ impl Capabilities for OpenRouterExt { type Rerank = Nothing; } -// OpenRouter's base URL already includes `/api/v1`, so the default -// `/embeddings` path is correct. impl crate::providers::openai::embedding::OpenAIEmbeddingsCompatible for OpenRouterExt {} impl DebugExt for OpenRouterExt {} diff --git a/crates/rig-core/src/providers/openrouter/embedding.rs b/crates/rig-core/src/providers/openrouter/embedding.rs index bff1ca71b..ce1f11d51 100644 --- a/crates/rig-core/src/providers/openrouter/embedding.rs +++ b/crates/rig-core/src/providers/openrouter/embedding.rs @@ -1,13 +1,42 @@ -//! OpenRouter's embeddings endpoint is OpenAI-compatible, so [`EmbeddingModel`] -//! is a thin alias over the shared [`GenericEmbeddingModel`]. OpenRouter's base -//! URL already carries the `/api/v1` segment, so the default `/embeddings` path -//! from the -//! [`OpenAIEmbeddingsCompatible`](crate::providers::openai::embedding::OpenAIEmbeddingsCompatible) -//! implementation on [`OpenRouterExt`](super::client::OpenRouterExt) is used. - use super::client::OpenRouterExt; use crate::providers::openai::embedding::GenericEmbeddingModel; -/// OpenRouter embedding model, backed by the shared OpenAI-compatible -/// [`GenericEmbeddingModel`]. pub type EmbeddingModel = GenericEmbeddingModel; + +#[cfg(test)] +mod tests { + use crate::client::EmbeddingsClient; + use crate::embeddings::EmbeddingModel as _; + use crate::providers::openrouter; + use crate::test_utils::RecordingHttpClient; + + const RESPONSE_BODY: &str = r#"{ + "id": "gen-1", + "object": "list", + "model": "openai/text-embedding-3-small", + "data": [{ "object": "embedding", "index": 0, "embedding": [0.5, 0.6] }] + }"#; + + #[tokio::test] + async fn openrouter_embeddings_use_default_path_and_zero_absent_usage() { + let http_client = RecordingHttpClient::new(RESPONSE_BODY); + let client = openrouter::Client::builder() + .api_key("dummy-key") + .http_client(http_client.clone()) + .build() + .expect("client should build"); + + let model = client.embedding_model("openai/text-embedding-3-small"); + let response = model + .embed_texts_with_usage(["hello".to_string()]) + .await + .expect("embedding request should succeed"); + + assert_eq!(response.embeddings.len(), 1); + assert_eq!(response.usage.total_tokens, 0); + + let requests = http_client.requests(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].uri, "https://openrouter.ai/api/v1/embeddings"); + } +} diff --git a/crates/rig-core/src/providers/together/client.rs b/crates/rig-core/src/providers/together/client.rs index f73f00085..3638302c2 100644 --- a/crates/rig-core/src/providers/together/client.rs +++ b/crates/rig-core/src/providers/together/client.rs @@ -48,7 +48,6 @@ impl crate::providers::openai::completion::OpenAICompatibleProvider for Together } impl crate::providers::openai::embedding::OpenAIEmbeddingsCompatible for TogetherExt { - // The client base URL is the bare host, so the version segment is explicit. fn embeddings_path(&self) -> String { "/v1/embeddings".to_string() } diff --git a/crates/rig-core/src/providers/together/embedding.rs b/crates/rig-core/src/providers/together/embedding.rs index 837228f71..2cb771ca4 100644 --- a/crates/rig-core/src/providers/together/embedding.rs +++ b/crates/rig-core/src/providers/together/embedding.rs @@ -3,12 +3,6 @@ //! From [Together AI Reference](https://docs.together.ai/docs/embeddings-overview) // ================================================================ -//! Together's embeddings endpoint is OpenAI-compatible, so [`EmbeddingModel`] is -//! a thin alias over the shared [`GenericEmbeddingModel`]. The `/v1/embeddings` -//! path is supplied by [`TogetherExt`](super::client::TogetherExt)'s -//! [`OpenAIEmbeddingsCompatible`](crate::providers::openai::embedding::OpenAIEmbeddingsCompatible) -//! implementation in `client.rs`. - use super::client::TogetherExt; use crate::providers::openai::embedding::GenericEmbeddingModel; @@ -25,9 +19,7 @@ pub const M2_BERT_80M_8K_RETRIEVAL: &str = "togethercomputer/m2-bert-80M-8k-retr pub const SENTENCE_BERT: &str = "sentence-transformers/msmarco-bert-base-dot-v5"; pub const UAE_LARGE_V1: &str = "WhereIsAI/UAE-Large-V1"; -/// Together AI embedding model, backed by the shared OpenAI-compatible -/// [`GenericEmbeddingModel`]. Unlike the previous hand-rolled implementation -/// this honors `ndims` (sent as `dimensions`), `encoding_format`, and `user`. +/// Together AI embedding model, driven by the shared OpenAI Embeddings path. pub type EmbeddingModel = GenericEmbeddingModel; #[cfg(test)] @@ -38,25 +30,25 @@ mod tests { use crate::providers::together; use crate::test_utils::RecordingHttpClient; - // Minimal OpenAI-compatible embeddings response. `usage` is intentionally - // omitted to exercise the optional-usage path the generic model now allows. const RESPONSE_BODY: &str = r#"{ "object": "list", "model": "BAAI/bge-base-en-v1.5", "data": [{ "object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3] }] }"#; + fn client(http_client: RecordingHttpClient) -> together::Client { + together::Client::builder() + .api_key("dummy-key") + .http_client(http_client) + .build() + .expect("client should build") + } + #[tokio::test] async fn together_embeddings_send_dimensions_to_v1_path() { let http_client = RecordingHttpClient::new(RESPONSE_BODY); - let client = together::Client::builder() - .api_key("dummy-key") - .http_client(http_client.clone()) - .build() - .expect("client should build"); + let client = client(http_client.clone()); - // With `ndims` set, `dimensions` must be sent — the previous hand-rolled - // Together model silently dropped it (the bug this migration fixes). let model = client.embedding_model_with_ndims(BGE_BASE_EN_V1_5, 3); model .embed_texts(["hello".to_string()]) @@ -75,4 +67,20 @@ mod tests { assert_eq!(body["dimensions"], serde_json::json!(3)); assert_eq!(body["model"], BGE_BASE_EN_V1_5); } + + #[tokio::test] + async fn together_embeddings_omit_dimensions_when_ndims_unset() { + let http_client = RecordingHttpClient::new(RESPONSE_BODY); + let client = client(http_client.clone()); + + let model = client.embedding_model(BGE_BASE_EN_V1_5); + model + .embed_texts(["hello".to_string()]) + .await + .expect("embedding request should succeed"); + + let body: serde_json::Value = serde_json::from_slice(&http_client.requests()[0].body) + .expect("request body should be JSON"); + assert!(body.get("dimensions").is_none()); + } }