From ca47b5b20bce85c797f56d99263c60446c319e72 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sun, 12 Jul 2026 11:20:49 +0000 Subject: [PATCH 1/4] feat(embeddings): port provider implementations --- Cargo.lock | 13 + Cargo.toml | 5 +- examples/local_model_probe.rs | 35 +-- src/harness/embeddings/cloud.rs | 121 ++++++++++ src/harness/embeddings/cohere.rs | 205 ++++++++++++++++ src/harness/embeddings/mod.rs | 21 ++ src/harness/embeddings/noop.rs | 29 +++ src/harness/embeddings/ollama.rs | 334 ++++++++++++++++++++++++++ src/harness/embeddings/openai.rs | 165 +++++++++++-- src/harness/embeddings/rate_limit.rs | 154 ++++++++++++ src/harness/embeddings/retry_after.rs | 67 ++++++ src/harness/embeddings/test.rs | 28 +++ src/harness/embeddings/types.rs | 26 ++ src/harness/embeddings/voyage.rs | 62 +++++ 14 files changed, 1231 insertions(+), 34 deletions(-) create mode 100644 src/harness/embeddings/cloud.rs create mode 100644 src/harness/embeddings/cohere.rs create mode 100644 src/harness/embeddings/noop.rs create mode 100644 src/harness/embeddings/ollama.rs create mode 100644 src/harness/embeddings/rate_limit.rs create mode 100644 src/harness/embeddings/retry_after.rs create mode 100644 src/harness/embeddings/voyage.rs diff --git a/Cargo.lock b/Cargo.lock index f323f5f..f3f6262 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1302,6 +1302,7 @@ dependencies = [ "sha2", "thiserror", "tokio", + "tracing", ] [[package]] @@ -1431,9 +1432,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "pin-project-lite", + "tracing-attributes", "tracing-core", ] +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tracing-core" version = "0.1.36" diff --git a/Cargo.toml b/Cargo.toml index b2d4d53..1d9df0d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,6 +25,7 @@ serde = { version = "1", features = ["derive"] } serde_json = "1" sha2 = "0.11" thiserror = "2" +tracing = "0.1" # Async runtime primitives used by the streaming pipeline, plus task spawning # for `DurabilityMode::Async` background checkpoint writes, `spawn_blocking` @@ -53,7 +54,7 @@ rusqlite = { version = "0.40", features = ["bundled"], optional = true } rhai = { version = "1", features = ["sync"], optional = true } # Optional builtin tool family for deterministic time/date helpers. -chrono = { version = "0.4", optional = true } +chrono = "0.4" chrono-tz = { version = "0.10", optional = true } [features] @@ -70,7 +71,7 @@ repl = ["dep:rhai"] rlm = ["dep:rhai", "tokio/process", "tokio/io-util"] # Builtin generic tools (`harness::tools`) kept out of the default dependency # graph so host applications can choose whether they want these implementations. -tools = ["dep:chrono", "dep:chrono-tz"] +tools = ["dep:chrono-tz"] [dev-dependencies] tokio = { version = "1", features = ["macros", "rt-multi-thread", "time", "test-util"] } diff --git a/examples/local_model_probe.rs b/examples/local_model_probe.rs index 48efd7a..84944bc 100644 --- a/examples/local_model_probe.rs +++ b/examples/local_model_probe.rs @@ -79,7 +79,6 @@ fn first_call(response: &ModelResponse) -> Option<&ToolCall> { response.message.tool_calls.first() } - fn msg_text(message: &AssistantMessage) -> String { message .content @@ -123,7 +122,10 @@ async fn tool_call(model: &OpenAiModel) -> Outcome { format!( "no tool_calls; finish={:?} text={:?}", resp.finish_reason, - msg_text(&resp.message).chars().take(120).collect::() + msg_text(&resp.message) + .chars() + .take(120) + .collect::() ), ), }, @@ -240,10 +242,7 @@ async fn parallel_tools(model: &OpenAiModel) -> Outcome { .iter() .map(|c| c.id.as_str()) .collect(); - let unique = ids - .iter() - .collect::>() - .len(); + let unique = ids.iter().collect::>().len(); if n >= 2 && unique == n { ok("parallel-tools", format!("{n} calls, ids={ids:?}")) } else if n >= 2 { @@ -251,7 +250,10 @@ async fn parallel_tools(model: &OpenAiModel) -> Outcome { } else { // Small models often serialize calls across turns; only flag // this as informational, not a harness bug. - ok("parallel-tools", format!("model made {n} call(s) (model behavior)")) + ok( + "parallel-tools", + format!("model made {n} call(s) (model behavior)"), + ) } } Err(e) => fail("parallel-tools", e.to_string()), @@ -313,16 +315,16 @@ async fn streaming_tools(model: &OpenAiModel) -> Outcome { "streaming-tools", format!("id={:?} args={}", call.id, call.arguments), ), - Some(call) => fail( - "streaming-tools", - format!("bad args: {}", call.arguments), - ), + Some(call) => fail("streaming-tools", format!("bad args: {}", call.arguments)), None => fail( "streaming-tools", format!( "no tool call; finish={:?} text={:?}", resp.finish_reason, - msg_text(&resp.message).chars().take(120).collect::() + msg_text(&resp.message) + .chars() + .take(120) + .collect::() ), ), }, @@ -345,7 +347,9 @@ async fn json_object(model: &OpenAiModel) -> Outcome { Ok(resp) => { let text = msg_text(&resp.message); match serde_json::from_str::(text.trim()) { - Ok(v) if v.is_object() => ok("json-object", text.chars().take(80).collect::()), + Ok(v) if v.is_object() => { + ok("json-object", text.chars().take(80).collect::()) + } _ => fail("json-object", format!("not a JSON object: {text:?}")), } } @@ -395,7 +399,10 @@ async fn thinking_leak(model: &OpenAiModel) -> Outcome { if text.contains("") || text.contains("") { fail( "thinking-leak", - format!(" leaked into text: {:?}", text.chars().take(120).collect::()), + format!( + " leaked into text: {:?}", + text.chars().take(120).collect::() + ), ) } else { ok("thinking-leak", text.chars().take(60).collect::()) diff --git a/src/harness/embeddings/cloud.rs b/src/harness/embeddings/cloud.rs new file mode 100644 index 0000000..b628f61 --- /dev/null +++ b/src/harness/embeddings/cloud.rs @@ -0,0 +1,121 @@ +//! Bearer-authenticated OpenAI-compatible cloud embedding model. + +use std::sync::Arc; + +use async_trait::async_trait; + +use super::{EmbeddingModel, OpenAiEmbeddingModel}; +use crate::error::{Result, TinyAgentsError}; + +pub const DEFAULT_CLOUD_MODEL: &str = "embedding-v1"; +pub const DEFAULT_CLOUD_DIMENSIONS: usize = 1024; + +/// Resolves the current bearer token for each request. +pub type BearerResolver = Arc Result + Send + Sync>; + +/// Cloud model whose credential lifecycle remains owned by the host. +pub struct CloudEmbeddingModel { + base_url: String, + model: String, + dimensions: usize, + bearer: BearerResolver, +} + +impl CloudEmbeddingModel { + pub fn new( + base_url: impl Into, + model: impl Into, + dimensions: usize, + bearer: BearerResolver, + ) -> Self { + Self { + base_url: base_url.into().trim().trim_end_matches('/').to_owned(), + model: model.into(), + dimensions, + bearer, + } + } +} + +#[async_trait] +impl EmbeddingModel for CloudEmbeddingModel { + fn name(&self) -> &str { + "cloud" + } + + fn model_id(&self) -> &str { + &self.model + } + + fn dimensions(&self) -> usize { + self.dimensions + } + + async fn embed(&self, texts: &[String]) -> Result>> { + if texts.is_empty() { + return Ok(Vec::new()); + } + if let Some(index) = texts.iter().position(|text| text.trim().is_empty()) { + return Err(TinyAgentsError::Validation(format!( + "cloud embed: refusing empty/whitespace input at index {index} of {} (model={})", + texts.len(), + self.model + ))); + } + let bearer = (self.bearer)()?; + if bearer.trim().is_empty() { + return Err(TinyAgentsError::Validation( + "No backend session for cloud embeddings".into(), + )); + } + OpenAiEmbeddingModel::new(bearer) + .with_base_url(&self.base_url) + .with_model(&self.model) + .with_dimensions(self.dimensions) + .with_send_dimensions(false) + .with_required_api_key(true) + .embed(texts) + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn missing_bearer() -> BearerResolver { + Arc::new(|| { + Err(TinyAgentsError::Validation( + "No backend session for cloud embeddings".into(), + )) + }) + } + + #[test] + fn identity_matches_host_contract() { + let model = CloudEmbeddingModel::new( + "https://api.example/openai/v1/", + DEFAULT_CLOUD_MODEL, + DEFAULT_CLOUD_DIMENSIONS, + missing_bearer(), + ); + assert_eq!(model.name(), "cloud"); + assert_eq!( + model.signature(), + "provider=cloud;model=embedding-v1;dims=1024" + ); + } + + #[tokio::test] + async fn validation_precedes_bearer_resolution() { + let model = CloudEmbeddingModel::new( + "https://api.example/openai/v1", + DEFAULT_CLOUD_MODEL, + DEFAULT_CLOUD_DIMENSIONS, + missing_bearer(), + ); + assert!(model.embed(&[]).await.unwrap().is_empty()); + let error = model.embed(&[" ".into()]).await.unwrap_err(); + assert!(error.to_string().contains("empty/whitespace")); + } +} diff --git a/src/harness/embeddings/cohere.rs b/src/harness/embeddings/cohere.rs new file mode 100644 index 0000000..4aa21b3 --- /dev/null +++ b/src/harness/embeddings/cohere.rs @@ -0,0 +1,205 @@ +//! Cohere-native embedding model using the `/v2/embed` contract. + +use async_trait::async_trait; +use serde::Deserialize; + +use super::EmbeddingModel; +use super::retry_after::{MAX_RETRIES, backoff_ms_for_attempt}; +use crate::error::{Result, TinyAgentsError}; + +pub const COHERE_API_BASE: &str = "https://api.cohere.com"; +pub const COHERE_DEFAULT_MODEL: &str = "embed-english-v3.0"; +pub const COHERE_DEFAULT_DIMENSIONS: usize = 1024; + +pub struct CohereEmbeddingModel { + client: reqwest::Client, + api_key: String, + model: String, + dimensions: usize, + base_url: String, +} + +impl CohereEmbeddingModel { + pub fn new(api_key: impl Into) -> Self { + Self { + client: reqwest::Client::new(), + api_key: api_key.into(), + model: COHERE_DEFAULT_MODEL.to_owned(), + dimensions: COHERE_DEFAULT_DIMENSIONS, + base_url: COHERE_API_BASE.to_owned(), + } + } + + pub fn with_model(mut self, model: impl Into) -> Self { + self.model = model.into(); + self + } + + pub fn with_dimensions(mut self, dimensions: usize) -> Self { + self.dimensions = dimensions; + self + } + + pub fn with_base_url(mut self, base_url: impl Into) -> Self { + self.base_url = base_url.into().trim().trim_end_matches('/').to_owned(); + self + } + + pub fn with_client(mut self, client: reqwest::Client) -> Self { + self.client = client; + self + } +} + +#[derive(Deserialize)] +struct CohereResponse { + embeddings: CohereEmbeddings, +} + +#[derive(Deserialize)] +struct CohereEmbeddings { + float: Vec>, +} + +#[async_trait] +impl EmbeddingModel for CohereEmbeddingModel { + fn name(&self) -> &str { + "cohere" + } + + fn model_id(&self) -> &str { + &self.model + } + + fn dimensions(&self) -> usize { + self.dimensions + } + + async fn embed(&self, texts: &[String]) -> Result>> { + if texts.is_empty() { + return Ok(Vec::new()); + } + if self.api_key.trim().is_empty() { + return Err(TinyAgentsError::Validation( + "Cohere API key not set. Configure an API key before embedding.".into(), + )); + } + if let Some(index) = texts.iter().position(|text| text.trim().is_empty()) { + return Err(TinyAgentsError::Validation(format!( + "cohere embed: refusing empty/whitespace input at index {index} of {}", + texts.len() + ))); + } + + let url = format!("{}/v2/embed", self.base_url); + let body = serde_json::json!({ + "model": self.model, + "texts": texts, + "input_type": "search_document", + "embedding_types": ["float"], + }); + + let mut response = None; + for attempt in 0..=MAX_RETRIES { + super::rate_limit::acquire(&self.base_url).await; + let current = self + .client + .post(&url) + .header("Authorization", format!("Bearer {}", self.api_key)) + .json(&body) + .send() + .await + .map_err(|error| { + TinyAgentsError::Embedding(format!( + "Cohere embeddings request to {url} failed: {error}" + )) + })?; + + let retryable = matches!(current.status().as_u16(), 429 | 503); + if retryable && attempt < MAX_RETRIES { + let retry_after = current + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + let status = current.status(); + let _ = current.text().await; + let delay_ms = backoff_ms_for_attempt(attempt, retry_after.as_deref()); + tracing::debug!( + target: "tinyagents::embeddings::cohere", + %status, + attempt, + delay_ms, + "[embeddings] retrying transient Cohere response" + ); + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + continue; + } + response = Some(current); + break; + } + + let response = response.expect("bounded retry loop always records its final response"); + let status = response.status(); + let text = response.text().await.map_err(|error| { + TinyAgentsError::Embedding(format!("Cohere embed response read failed: {error}")) + })?; + if !status.is_success() { + return Err(TinyAgentsError::Embedding(format!( + "Cohere embed API error ({status}): {text}" + ))); + } + + let payload: CohereResponse = serde_json::from_str(&text).map_err(|error| { + TinyAgentsError::Embedding(format!("Cohere embed response parse failed: {error}")) + })?; + let vectors = payload.embeddings.float; + if vectors.len() != texts.len() { + return Err(TinyAgentsError::Embedding(format!( + "Cohere embed count mismatch: sent {} texts, got {} embeddings", + texts.len(), + vectors.len() + ))); + } + for (index, vector) in vectors.iter().enumerate() { + if self.dimensions > 0 && vector.len() != self.dimensions { + return Err(TinyAgentsError::Embedding(format!( + "Cohere embed dimension mismatch at index {index}: expected {}, got {}", + self.dimensions, + vector.len() + ))); + } + } + Ok(vectors) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn identity_and_defaults_match_host_contract() { + let model = CohereEmbeddingModel::new("key"); + assert_eq!(model.name(), "cohere"); + assert_eq!(model.model_id(), COHERE_DEFAULT_MODEL); + assert_eq!(model.dimensions(), COHERE_DEFAULT_DIMENSIONS); + assert_eq!( + model.signature(), + "provider=cohere;model=embed-english-v3.0;dims=1024" + ); + } + + #[tokio::test] + async fn empty_batch_short_circuits_before_key_validation() { + let model = CohereEmbeddingModel::new(""); + assert!(model.embed(&[]).await.unwrap().is_empty()); + } + + #[tokio::test] + async fn missing_key_fails_before_network() { + let model = CohereEmbeddingModel::new("").with_base_url("http://127.0.0.1:1"); + let error = model.embed(&["hello".into()]).await.unwrap_err(); + assert!(error.to_string().contains("API key not set")); + } +} diff --git a/src/harness/embeddings/mod.rs b/src/harness/embeddings/mod.rs index 089cd03..85269ab 100644 --- a/src/harness/embeddings/mod.rs +++ b/src/harness/embeddings/mod.rs @@ -273,9 +273,30 @@ impl Retriever { } } +mod cloud; +mod cohere; +mod noop; +mod ollama; mod openai; +pub mod rate_limit; +pub mod retry_after; +mod voyage; +pub use noop::NoopEmbeddingModel; +pub use ollama::{ + DEFAULT_OLLAMA_DIMENSIONS, DEFAULT_OLLAMA_MODEL, DEFAULT_OLLAMA_URL, OllamaEmbeddingModel, +}; pub use openai::OpenAiEmbeddingModel; +pub use types::format_embedding_signature; +pub use voyage::{ + VOYAGE_API_BASE, VOYAGE_DEFAULT_DIMENSIONS, VOYAGE_DEFAULT_MODEL, VoyageEmbeddingModel, +}; #[cfg(test)] mod test; +pub use cloud::{ + BearerResolver, CloudEmbeddingModel, DEFAULT_CLOUD_DIMENSIONS, DEFAULT_CLOUD_MODEL, +}; +pub use cohere::{ + COHERE_API_BASE, COHERE_DEFAULT_DIMENSIONS, COHERE_DEFAULT_MODEL, CohereEmbeddingModel, +}; diff --git a/src/harness/embeddings/noop.rs b/src/harness/embeddings/noop.rs new file mode 100644 index 0000000..2133319 --- /dev/null +++ b/src/harness/embeddings/noop.rs @@ -0,0 +1,29 @@ +//! No-op embedding model for keyword-only retrieval. + +use async_trait::async_trait; + +use super::EmbeddingModel; +use crate::error::Result; + +/// Embedding model used when semantic search is disabled. +#[derive(Clone, Copy, Debug, Default)] +pub struct NoopEmbeddingModel; + +#[async_trait] +impl EmbeddingModel for NoopEmbeddingModel { + fn name(&self) -> &str { + "none" + } + + fn model_id(&self) -> &str { + "none" + } + + fn dimensions(&self) -> usize { + 0 + } + + async fn embed(&self, _texts: &[String]) -> Result>> { + Ok(Vec::new()) + } +} diff --git a/src/harness/embeddings/ollama.rs b/src/harness/embeddings/ollama.rs new file mode 100644 index 0000000..96ad6cb --- /dev/null +++ b/src/harness/embeddings/ollama.rs @@ -0,0 +1,334 @@ +//! Ollama `/api/embed` model with positional and NaN recovery guarantees. + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; + +use super::EmbeddingModel; +use crate::error::{Result, TinyAgentsError}; + +pub const DEFAULT_OLLAMA_URL: &str = "http://localhost:11434"; +pub const DEFAULT_OLLAMA_MODEL: &str = "bge-m3"; +pub const DEFAULT_OLLAMA_DIMENSIONS: usize = 1024; + +#[derive(Debug)] +pub struct OllamaEmbeddingModel { + client: reqwest::Client, + base_url: String, + model: String, + dimensions: usize, +} + +impl OllamaEmbeddingModel { + pub fn try_new(base_url: &str, model: &str, dimensions: usize) -> Result { + Ok(Self { + client: reqwest::Client::new(), + base_url: normalize_base_url(base_url)?, + model: normalize_model(model)?, + dimensions: if dimensions == 0 { + DEFAULT_OLLAMA_DIMENSIONS + } else { + dimensions + }, + }) + } + + pub fn new(base_url: &str, model: &str, dimensions: usize) -> Self { + Self::try_new(base_url, model, dimensions).expect("invalid Ollama embedding configuration") + } + + pub fn with_client(mut self, client: reqwest::Client) -> Self { + self.client = client; + self + } + + pub fn base_url(&self) -> &str { + &self.base_url + } + + pub fn model(&self) -> &str { + &self.model + } + + pub fn embed_url(&self) -> String { + format!("{}/api/embed", self.base_url) + } + + async fn request(&self, input: Vec) -> Result { + self.client + .post(self.embed_url()) + .json(&OllamaRequest { + model: self.model.clone(), + input, + }) + .send() + .await + .map_err(|error| { + TinyAgentsError::Embedding(format!( + "ollama embed request failed (is Ollama running at {}?): {error}", + self.base_url + )) + }) + } + + async fn embed_one_with_nan_recovery(&self, text: &str) -> Result>> { + let response = self.request(vec![text.to_owned()]).await?; + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + if status.as_u16() == 500 && is_nan_encode_error(&body) { + tracing::warn!( + target: "tinyagents::embeddings::ollama", + model = self.model, + "[embeddings] substituting empty vector for Ollama NaN response" + ); + return Ok(None); + } + return Err(ollama_http_error(status, &body)); + } + let payload = parse_response(response).await?; + if payload.embeddings.len() != 1 { + return Err(TinyAgentsError::Embedding(format!( + "ollama embed count mismatch: sent 1 text, got {} embeddings", + payload.embeddings.len() + ))); + } + let vector = payload.embeddings.into_iter().next().unwrap(); + self.validate_dimensions(0, &vector)?; + Ok(Some(vector)) + } + + async fn embed_per_text( + &self, + total: usize, + live: &[(usize, String)], + ) -> Result>> { + let mut output = vec![Vec::new(); total]; + for (index, text) in live { + if let Some(vector) = self.embed_one_with_nan_recovery(text).await? { + output[*index] = vector; + } + } + Ok(output) + } + + fn validate_dimensions(&self, index: usize, vector: &[f32]) -> Result<()> { + if vector.len() != self.dimensions { + return Err(TinyAgentsError::Embedding(format!( + "ollama embed dimension mismatch at index {index}: expected {}, got {}", + self.dimensions, + vector.len() + ))); + } + Ok(()) + } +} + +impl Default for OllamaEmbeddingModel { + fn default() -> Self { + Self::new( + DEFAULT_OLLAMA_URL, + DEFAULT_OLLAMA_MODEL, + DEFAULT_OLLAMA_DIMENSIONS, + ) + } +} + +#[derive(Serialize)] +struct OllamaRequest { + model: String, + input: Vec, +} + +#[derive(Deserialize)] +struct OllamaResponse { + #[serde(default)] + embeddings: Vec>, +} + +async fn parse_response(response: reqwest::Response) -> Result { + response.json().await.map_err(|error| { + TinyAgentsError::Embedding(format!("ollama embed response parse failed: {error}")) + }) +} + +fn ollama_http_error(status: reqwest::StatusCode, body: &str) -> TinyAgentsError { + let detail = body.trim(); + TinyAgentsError::Embedding(format!( + "ollama embed failed with status {status}{}", + if detail.is_empty() { + String::new() + } else { + format!(": {detail}") + } + )) +} + +fn is_nan_encode_error(body: &str) -> bool { + body.to_ascii_lowercase().contains("unsupported value: nan") +} + +fn normalize_base_url(base_url: &str) -> Result { + let raw = if base_url.trim().is_empty() { + DEFAULT_OLLAMA_URL + } else { + base_url.trim() + }; + let url = reqwest::Url::parse(raw).map_err(|error| { + TinyAgentsError::Validation(format!("invalid Ollama base_url `{raw}`: {error}")) + })?; + if !matches!(url.scheme(), "http" | "https") { + return Err(TinyAgentsError::Validation(format!( + "invalid Ollama base_url `{raw}`: expected an http:// or https:// URL" + ))); + } + if !url.username().is_empty() || url.password().is_some() { + return Err(TinyAgentsError::Validation(format!( + "invalid Ollama base_url `{raw}`: configure the server root without credentials" + ))); + } + if url.query().is_some() || url.fragment().is_some() { + return Err(TinyAgentsError::Validation(format!( + "invalid Ollama base_url `{raw}`: query strings and fragments are not supported" + ))); + } + let segments = url + .path_segments() + .map(|parts| { + parts + .filter(|part| !part.is_empty()) + .map(str::to_ascii_lowercase) + .collect::>() + }) + .unwrap_or_default(); + let endpoint_suffix = segments + .iter() + .any(|part| matches!(part.as_str(), "api" | "v1")) + || (segments.len() >= 2 + && segments[segments.len() - 2] == "chat" + && segments[segments.len() - 1] == "completions"); + if endpoint_suffix { + return Err(TinyAgentsError::Validation(format!( + "invalid Ollama base_url `{raw}`: configure the Ollama server root, not an API endpoint" + ))); + } + Ok(url.as_str().trim_end_matches('/').to_owned()) +} + +fn normalize_model(model: &str) -> Result { + let model = if model.trim().is_empty() { + DEFAULT_OLLAMA_MODEL.to_owned() + } else { + model.trim().to_owned() + }; + if model.to_ascii_lowercase().starts_with("local-") { + return Err(TinyAgentsError::Validation(format!( + "invalid Ollama embedding model `{model}`: `local-*` IDs are virtual routing aliases; configure `{DEFAULT_OLLAMA_MODEL}` or another real model" + ))); + } + Ok(model) +} + +#[async_trait] +impl EmbeddingModel for OllamaEmbeddingModel { + fn name(&self) -> &str { + "ollama" + } + + fn model_id(&self) -> &str { + &self.model + } + + fn dimensions(&self) -> usize { + self.dimensions + } + + async fn embed(&self, texts: &[String]) -> Result>> { + if texts.is_empty() { + return Ok(Vec::new()); + } + let live = texts + .iter() + .enumerate() + .filter_map(|(index, text)| { + let text = text.trim(); + (!text.is_empty()).then(|| (index, text.to_owned())) + }) + .collect::>(); + if live.is_empty() { + return Ok(vec![Vec::new(); texts.len()]); + } + let input = live + .iter() + .map(|(_, text)| text.clone()) + .collect::>(); + let response = self.request(input.clone()).await?; + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + if status.as_u16() == 500 && is_nan_encode_error(&body) { + tracing::warn!( + target: "tinyagents::embeddings::ollama", + batch = live.len(), + model = self.model, + "[embeddings] recovering Ollama NaN batch per text" + ); + if live.len() == 1 { + return Ok(vec![Vec::new(); texts.len()]); + } + return self.embed_per_text(texts.len(), &live).await; + } + return Err(ollama_http_error(status, &body)); + } + + let payload = parse_response(response).await?; + if payload.embeddings.len() != input.len() { + return Err(TinyAgentsError::Embedding(format!( + "ollama embed count mismatch: sent {} texts, got {} embeddings", + input.len(), + payload.embeddings.len() + ))); + } + for (index, vector) in payload.embeddings.iter().enumerate() { + self.validate_dimensions(index, vector)?; + } + let mut output = vec![Vec::new(); texts.len()]; + for ((index, _), vector) in live.iter().zip(payload.embeddings) { + output[*index] = vector; + } + Ok(output) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn defaults_and_identity_match_host() { + let model = OllamaEmbeddingModel::default(); + assert_eq!(model.base_url(), DEFAULT_OLLAMA_URL); + assert_eq!(model.model_id(), DEFAULT_OLLAMA_MODEL); + assert_eq!(model.dimensions(), DEFAULT_OLLAMA_DIMENSIONS); + assert_eq!(model.signature(), "provider=ollama;model=bge-m3;dims=1024"); + } + + #[test] + fn validates_root_url_and_real_model() { + assert!(OllamaEmbeddingModel::try_new("http://host:11434/api", "m", 1).is_err()); + assert!(OllamaEmbeddingModel::try_new("http://user:p@host:11434", "m", 1).is_err()); + assert!(OllamaEmbeddingModel::try_new("http://host:11434", "local-v1", 1).is_err()); + } + + #[tokio::test] + async fn blank_inputs_preserve_positions_without_network() { + let model = OllamaEmbeddingModel::default(); + let vectors = model.embed(&[" ".into(), "\n".into()]).await.unwrap(); + assert_eq!(vectors, vec![Vec::::new(), Vec::new()]); + } + + #[test] + fn recognizes_only_nan_encoding_failures() { + assert!(is_nan_encode_error("unsupported value: NaN")); + assert!(!is_nan_encode_error("model crashed")); + } +} diff --git a/src/harness/embeddings/openai.rs b/src/harness/embeddings/openai.rs index 8316965..1fab0da 100644 --- a/src/harness/embeddings/openai.rs +++ b/src/harness/embeddings/openai.rs @@ -8,6 +8,7 @@ use async_trait::async_trait; use serde_json::{Value, json}; use super::EmbeddingModel; +use super::retry_after::{MAX_RETRIES, backoff_ms_for_attempt}; use crate::error::{Result, TinyAgentsError}; /// Default OpenAI embedding model id. @@ -43,6 +44,8 @@ pub struct OpenAiEmbeddingModel { model: String, base_url: String, dimensions: usize, + send_dimensions: bool, + requires_api_key: bool, } impl OpenAiEmbeddingModel { @@ -55,12 +58,14 @@ impl OpenAiEmbeddingModel { model: DEFAULT_MODEL.to_string(), base_url: DEFAULT_BASE_URL.to_string(), dimensions: DEFAULT_DIMENSIONS, + send_dimensions: true, + requires_api_key: true, } } /// Overrides the embedding model id. pub fn with_model(mut self, model: impl Into) -> Self { - self.model = model.into(); + self.model = normalize_model_for_base_url(&self.base_url, &model.into()); self } @@ -68,6 +73,7 @@ impl OpenAiEmbeddingModel { /// endpoint is always `{base_url}/embeddings`. pub fn with_base_url(mut self, base_url: impl Into) -> Self { self.base_url = base_url.into().trim_end_matches('/').to_string(); + self.model = normalize_model_for_base_url(&self.base_url, &self.model); self } @@ -78,6 +84,30 @@ impl OpenAiEmbeddingModel { self } + /// Controls whether the OpenAI-compatible `dimensions` field is sent. + pub fn with_send_dimensions(mut self, send: bool) -> Self { + self.send_dimensions = send; + self + } + + /// Controls whether an empty API key is rejected before making a request. + pub fn with_required_api_key(mut self, required: bool) -> Self { + self.requires_api_key = required; + self + } + + pub fn base_url(&self) -> &str { + &self.base_url + } + + pub fn model(&self) -> &str { + &self.model + } + + pub fn embeddings_url(&self) -> String { + embeddings_url(&self.base_url) + } + /// Builds a model from environment variables. /// /// Reads `OPENAI_API_KEY` (required), `OPENAI_EMBEDDING_MODEL` @@ -112,29 +142,76 @@ impl OpenAiEmbeddingModel { #[async_trait] impl EmbeddingModel for OpenAiEmbeddingModel { + fn name(&self) -> &str { + "openai" + } + + fn model_id(&self) -> &str { + &self.model + } + async fn embed(&self, texts: &[String]) -> Result>> { if texts.is_empty() { return Ok(Vec::new()); } - let url = format!("{}/embeddings", self.base_url); - let body = json!({ + if let Some(index) = texts.iter().position(|text| text.trim().is_empty()) { + return Err(TinyAgentsError::Validation(format!( + "openai embed: refusing empty/whitespace input at index {index} of {} (model={})", + texts.len(), + self.model + ))); + } + if self.requires_api_key && self.api_key.trim().is_empty() { + return Err(TinyAgentsError::Validation(format!( + "Embedding API key not set (model={})", + self.model + ))); + } + let url = self.embeddings_url(); + let mut body = json!({ "model": self.model, "input": texts, - "dimensions": self.dimensions, }); + if self.send_dimensions && self.dimensions > 0 { + body["dimensions"] = json!(self.dimensions); + } - let response = self - .client - .post(&url) - .header("Authorization", format!("Bearer {}", self.api_key)) - .json(&body) - .send() - .await - .map_err(|e| { + let mut response = None; + for attempt in 0..=MAX_RETRIES { + super::rate_limit::acquire(&self.base_url).await; + let mut request = self.client.post(&url).json(&body); + if !self.api_key.is_empty() { + request = request.header("Authorization", format!("Bearer {}", self.api_key)); + } + let current = request.send().await.map_err(|e| { TinyAgentsError::Embedding(format!( "openai embeddings request to {url} failed: {e}" )) })?; + let retryable = matches!(current.status().as_u16(), 429 | 503); + if retryable && attempt < MAX_RETRIES { + let retry_after = current + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + let status = current.status(); + let _ = current.text().await; + let delay_ms = backoff_ms_for_attempt(attempt, retry_after.as_deref()); + tracing::debug!( + target: "tinyagents::embeddings::openai", + %status, + attempt, + delay_ms, + "[embeddings] retrying transient OpenAI-compatible response" + ); + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + continue; + } + response = Some(current); + break; + } + let response = response.expect("bounded retry loop always records its final response"); let status = response.status(); let text = response.text().await.map_err(|e| { @@ -163,12 +240,31 @@ impl EmbeddingModel for OpenAiEmbeddingModel { "openai embeddings response missing `embedding` array".into(), ) })?; - vectors.push( - embedding - .iter() - .map(|n| n.as_f64().unwrap_or(0.0) as f32) - .collect(), - ); + let vector = embedding + .iter() + .map(|n| { + n.as_f64().map(|value| value as f32).ok_or_else(|| { + TinyAgentsError::Embedding( + "openai embeddings response contains a non-numeric value".into(), + ) + }) + }) + .collect::>>()?; + if self.dimensions > 0 && vector.len() != self.dimensions { + return Err(TinyAgentsError::Embedding(format!( + "openai embed dimension mismatch: expected {}, got {}", + self.dimensions, + vector.len() + ))); + } + vectors.push(vector); + } + if vectors.len() != texts.len() { + return Err(TinyAgentsError::Embedding(format!( + "openai embed count mismatch: sent {} texts, got {} embeddings", + texts.len(), + vectors.len() + ))); } Ok(vectors) } @@ -177,3 +273,36 @@ impl EmbeddingModel for OpenAiEmbeddingModel { self.dimensions } } + +fn embeddings_url(base_url: &str) -> String { + let Ok(url) = reqwest::Url::parse(base_url) else { + return format!("{}/v1/embeddings", base_url.trim_end_matches('/')); + }; + let path = url.path().trim_end_matches('/'); + if path.ends_with("/embeddings") { + base_url.trim_end_matches('/').to_owned() + } else if path.is_empty() || path == "/" { + format!("{}/v1/embeddings", base_url.trim_end_matches('/')) + } else { + format!("{}/embeddings", base_url.trim_end_matches('/')) + } +} + +fn normalize_model_for_base_url(base_url: &str, model: &str) -> String { + let is_gemini = reqwest::Url::parse(base_url) + .ok() + .and_then(|url| url.host_str().map(str::to_ascii_lowercase)) + .is_some_and(|host| { + host == "generativelanguage.googleapis.com" + || host.ends_with(".generativelanguage.googleapis.com") + }); + if is_gemini + && !model.is_empty() + && !model.starts_with("models/") + && !model.starts_with("tunedModels/") + { + format!("models/{model}") + } else { + model.to_owned() + } +} diff --git a/src/harness/embeddings/rate_limit.rs b/src/harness/embeddings/rate_limit.rs new file mode 100644 index 0000000..4e90dfd --- /dev/null +++ b/src/harness/embeddings/rate_limit.rs @@ -0,0 +1,154 @@ +//! Process-wide endpoint-keyed pacing for outbound embedding requests. + +use std::collections::HashMap; +use std::net::IpAddr; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::{Arc, Mutex, OnceLock, PoisonError}; +use std::time::Duration; + +use tokio::time::Instant; + +pub const DEFAULT_REQUESTS_PER_MINUTE: u32 = 60; + +static CONFIGURED_LIMIT: AtomicU32 = AtomicU32::new(DEFAULT_REQUESTS_PER_MINUTE); +static BUCKETS: OnceLock>>> = OnceLock::new(); + +pub fn set_rate_limit(per_minute: u32) { + let previous = CONFIGURED_LIMIT.swap(per_minute, Ordering::Relaxed); + if previous != per_minute { + if let Some(registry) = BUCKETS.get() { + registry + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clear(); + } + } + tracing::debug!( + target: "tinyagents::embeddings::rate_limit", + per_minute, + "[embeddings] configured outbound request rate" + ); +} + +pub fn rate_limit() -> u32 { + CONFIGURED_LIMIT.load(Ordering::Relaxed) +} + +pub async fn acquire(base_url: &str) { + acquire_with_limit(base_url, rate_limit()).await; +} + +async fn acquire_with_limit(base_url: &str, limit: u32) { + if limit == 0 || is_loopback_url(base_url) { + return; + } + bucket_for(base_url, limit).acquire().await; +} + +fn bucket_for(base_url: &str, per_minute: u32) -> Arc { + let registry = BUCKETS.get_or_init(|| Mutex::new(HashMap::new())); + let mut buckets = registry.lock().unwrap_or_else(PoisonError::into_inner); + buckets + .entry(base_url.to_owned()) + .or_insert_with(|| Arc::new(TokenBucket::per_minute(per_minute))) + .clone() +} + +fn is_loopback_url(base_url: &str) -> bool { + let Ok(url) = reqwest::Url::parse(base_url) else { + return false; + }; + url.host_str().is_some_and(|host| { + host.eq_ignore_ascii_case("localhost") + || host + .trim_start_matches('[') + .trim_end_matches(']') + .parse::() + .is_ok_and(|ip| ip.is_loopback()) + }) +} + +struct TokenBucket { + state: tokio::sync::Mutex, + refill_per_second: f64, +} + +struct BucketState { + tokens: f64, + last_refill: Instant, +} + +impl TokenBucket { + fn per_minute(per_minute: u32) -> Self { + Self { + state: tokio::sync::Mutex::new(BucketState { + tokens: 1.0, + last_refill: Instant::now(), + }), + refill_per_second: f64::from(per_minute.max(1)) / 60.0, + } + } + + async fn acquire(&self) { + loop { + let wait = { + let mut state = self.state.lock().await; + let now = Instant::now(); + let elapsed = now.duration_since(state.last_refill).as_secs_f64(); + state.last_refill = now; + refill_and_take(&mut state.tokens, self.refill_per_second, elapsed) + }; + let Some(wait) = wait else { + return; + }; + tracing::debug!( + target: "tinyagents::embeddings::rate_limit", + wait_ms = wait.as_millis(), + "[embeddings] waiting for outbound request slot" + ); + tokio::time::sleep(wait).await; + } + } +} + +fn refill_and_take( + tokens: &mut f64, + refill_per_second: f64, + elapsed_seconds: f64, +) -> Option { + *tokens = (*tokens + elapsed_seconds * refill_per_second).min(1.0); + if *tokens >= 1.0 { + *tokens -= 1.0; + None + } else { + Some(Duration::from_secs_f64((1.0 - *tokens) / refill_per_second)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn loopback_detection_is_fail_closed() { + assert!(is_loopback_url("http://localhost:11434")); + assert!(is_loopback_url("http://127.0.0.1:8080")); + assert!(is_loopback_url("http://[::1]:8080")); + assert!(!is_loopback_url("https://api.openai.com")); + assert!(!is_loopback_url("not a url")); + } + + #[test] + fn bucket_math_paces_without_bursting() { + let mut tokens = 1.0; + assert!(refill_and_take(&mut tokens, 1.0, 0.0).is_none()); + let wait = refill_and_take(&mut tokens, 1.0, 0.25).unwrap(); + assert!((wait.as_secs_f64() - 0.75).abs() < 1e-6); + } + + #[tokio::test] + async fn disabled_and_loopback_limits_never_block() { + acquire_with_limit("https://api.example.com", 0).await; + acquire_with_limit("http://127.0.0.1:1", 1).await; + } +} diff --git a/src/harness/embeddings/retry_after.rs b/src/harness/embeddings/retry_after.rs new file mode 100644 index 0000000..cdb5d37 --- /dev/null +++ b/src/harness/embeddings/retry_after.rs @@ -0,0 +1,67 @@ +//! Retry-After parsing and bounded exponential backoff for embedding providers. + +use chrono::{DateTime, Utc}; + +pub const MAX_RETRIES: u32 = 3; +pub const BASE_BACKOFF_MS: u64 = 1_000; +pub const MAX_BACKOFF_MS: u64 = 30_000; + +pub fn parse_retry_after_ms(value: Option<&str>) -> Option { + parse_retry_after_ms_at(value, Utc::now()) +} + +fn parse_retry_after_ms_at(value: Option<&str>, now: DateTime) -> Option { + let value = value?.trim(); + if value.is_empty() { + return None; + } + if let Ok(seconds) = value.parse::() { + return Some(seconds.saturating_mul(1_000).min(MAX_BACKOFF_MS)); + } + let retry_at = DateTime::parse_from_rfc2822(value) + .ok()? + .with_timezone(&Utc); + let delay = retry_at.signed_duration_since(now).num_milliseconds(); + if delay <= 0 { + return Some(0); + } + u64::try_from(delay).ok().map(|ms| ms.min(MAX_BACKOFF_MS)) +} + +pub fn backoff_ms_for_attempt(attempt: u32, retry_after: Option<&str>) -> u64 { + parse_retry_after_ms(retry_after).unwrap_or_else(|| { + BASE_BACKOFF_MS + .saturating_mul(2u64.saturating_pow(attempt)) + .min(MAX_BACKOFF_MS) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_delta_seconds_and_caps() { + assert_eq!(parse_retry_after_ms(Some(" 5 ")), Some(5_000)); + assert_eq!(parse_retry_after_ms(Some("999")), Some(MAX_BACKOFF_MS)); + assert_eq!(parse_retry_after_ms(Some("-1")), None); + } + + #[test] + fn parses_http_dates() { + let now = DateTime::parse_from_rfc2822("Wed, 21 Oct 2015 07:27:55 GMT") + .unwrap() + .with_timezone(&Utc); + assert_eq!( + parse_retry_after_ms_at(Some("Wed, 21 Oct 2015 07:28:00 GMT"), now), + Some(5_000) + ); + } + + #[test] + fn falls_back_to_bounded_exponential_backoff() { + assert_eq!(backoff_ms_for_attempt(0, None), 1_000); + assert_eq!(backoff_ms_for_attempt(2, None), 4_000); + assert_eq!(backoff_ms_for_attempt(20, None), MAX_BACKOFF_MS); + } +} diff --git a/src/harness/embeddings/test.rs b/src/harness/embeddings/test.rs index 8a23650..f6555bb 100644 --- a/src/harness/embeddings/test.rs +++ b/src/harness/embeddings/test.rs @@ -255,3 +255,31 @@ async fn retriever_accessors_expose_collaborators() { assert_eq!(retriever.model().dimensions(), 8); let _ = retriever.store(); } +#[test] +fn embedding_identity_signature_is_stable() { + let model = MockEmbeddingModel::new(8); + assert_eq!(model.name(), "mock"); + assert_eq!(model.model_id(), "deterministic-hash"); + assert_eq!( + model.signature(), + "provider=mock;model=deterministic-hash;dims=8" + ); + assert_eq!( + format_embedding_signature("openai", "text-embedding-3-small", 1536), + "provider=openai;model=text-embedding-3-small;dims=1536" + ); +} + +#[test] +fn voyage_and_noop_identity_match_host_contract() { + let voyage = VoyageEmbeddingModel::new("test-key"); + assert_eq!(voyage.name(), "voyage"); + assert_eq!(voyage.model_id(), VOYAGE_DEFAULT_MODEL); + assert_eq!( + voyage.signature(), + "provider=voyage;model=voyage-3-large;dims=1024" + ); + + let noop = NoopEmbeddingModel; + assert_eq!(noop.signature(), "provider=none;model=none;dims=0"); +} diff --git a/src/harness/embeddings/types.rs b/src/harness/embeddings/types.rs index 2c38d03..6b2b138 100644 --- a/src/harness/embeddings/types.rs +++ b/src/harness/embeddings/types.rs @@ -39,6 +39,12 @@ use crate::error::Result; /// deterministic implementations such as [`MockEmbeddingModel`]. #[async_trait] pub trait EmbeddingModel: Send + Sync { + /// Stable provider identifier, such as `"openai"` or `"ollama"`. + fn name(&self) -> &str; + + /// Stable model identifier within the provider. + fn model_id(&self) -> &str; + /// Embeds a batch of texts, returning one vector per input in input order. /// /// Returning an empty `Vec` for empty input is valid. @@ -46,6 +52,18 @@ pub trait EmbeddingModel: Send + Sync { /// Returns the fixed dimensionality of every vector this model produces. fn dimensions(&self) -> usize; + + /// Stable embedding-space identity used to partition persisted vectors. + fn signature(&self) -> String { + format_embedding_signature(self.name(), self.model_id(), self.dimensions()) + } +} + +/// Format the canonical embedding-space signature. +/// +/// This must remain byte-identical to OpenHuman's persisted signature contract. +pub fn format_embedding_signature(name: &str, model_id: &str, dimensions: usize) -> String { + format!("provider={name};model={model_id};dims={dimensions}") } // ── MockEmbeddingModel ──────────────────────────────────────────────────────── @@ -115,6 +133,14 @@ impl MockEmbeddingModel { #[async_trait] impl EmbeddingModel for MockEmbeddingModel { + fn name(&self) -> &str { + "mock" + } + + fn model_id(&self) -> &str { + "deterministic-hash" + } + async fn embed(&self, texts: &[String]) -> Result>> { Ok(texts.iter().map(|t| self.embed_one(t)).collect()) } diff --git a/src/harness/embeddings/voyage.rs b/src/harness/embeddings/voyage.rs new file mode 100644 index 0000000..bf488bc --- /dev/null +++ b/src/harness/embeddings/voyage.rs @@ -0,0 +1,62 @@ +//! Voyage AI embedding model. + +use async_trait::async_trait; + +use super::{EmbeddingModel, OpenAiEmbeddingModel}; +use crate::error::Result; + +pub const VOYAGE_API_BASE: &str = "https://api.voyageai.com/v1"; +pub const VOYAGE_DEFAULT_MODEL: &str = "voyage-3-large"; +pub const VOYAGE_DEFAULT_DIMENSIONS: usize = 1024; + +/// Voyage's endpoint uses the OpenAI response shape without its `dimensions` +/// request parameter. +pub struct VoyageEmbeddingModel { + inner: OpenAiEmbeddingModel, +} + +impl VoyageEmbeddingModel { + pub fn new(api_key: impl Into) -> Self { + Self::with_options( + api_key, + VOYAGE_DEFAULT_MODEL, + VOYAGE_DEFAULT_DIMENSIONS, + VOYAGE_API_BASE, + ) + } + + pub fn with_options( + api_key: impl Into, + model: impl Into, + dimensions: usize, + base_url: impl Into, + ) -> Self { + Self { + inner: OpenAiEmbeddingModel::new(api_key) + .with_model(model) + .with_dimensions(dimensions) + .with_base_url(base_url) + .with_send_dimensions(false) + .with_required_api_key(true), + } + } +} + +#[async_trait] +impl EmbeddingModel for VoyageEmbeddingModel { + fn name(&self) -> &str { + "voyage" + } + + fn model_id(&self) -> &str { + self.inner.model_id() + } + + fn dimensions(&self) -> usize { + self.inner.dimensions() + } + + async fn embed(&self, texts: &[String]) -> Result>> { + self.inner.embed(texts).await + } +} From f60981bdce345153c2d19a490a2a54fbe5168246 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sun, 12 Jul 2026 11:37:34 +0000 Subject: [PATCH 2/4] fix(embeddings): preserve provider contracts --- src/harness/embeddings/cloud.rs | 3 +++ src/harness/embeddings/mod.rs | 8 ++++++-- src/harness/embeddings/noop.rs | 4 ++-- src/harness/embeddings/openai.rs | 27 ++++++--------------------- src/harness/embeddings/rate_limit.rs | 14 +++++++------- src/harness/embeddings/test.rs | 18 ++++++++++++++++-- 6 files changed, 40 insertions(+), 34 deletions(-) diff --git a/src/harness/embeddings/cloud.rs b/src/harness/embeddings/cloud.rs index b628f61..0c2c664 100644 --- a/src/harness/embeddings/cloud.rs +++ b/src/harness/embeddings/cloud.rs @@ -15,6 +15,7 @@ pub type BearerResolver = Arc Result + Send + Sync>; /// Cloud model whose credential lifecycle remains owned by the host. pub struct CloudEmbeddingModel { + client: reqwest::Client, base_url: String, model: String, dimensions: usize, @@ -29,6 +30,7 @@ impl CloudEmbeddingModel { bearer: BearerResolver, ) -> Self { Self { + client: reqwest::Client::new(), base_url: base_url.into().trim().trim_end_matches('/').to_owned(), model: model.into(), dimensions, @@ -69,6 +71,7 @@ impl EmbeddingModel for CloudEmbeddingModel { )); } OpenAiEmbeddingModel::new(bearer) + .with_client(self.client.clone()) .with_base_url(&self.base_url) .with_model(&self.model) .with_dimensions(self.dimensions) diff --git a/src/harness/embeddings/mod.rs b/src/harness/embeddings/mod.rs index 85269ab..2b54448 100644 --- a/src/harness/embeddings/mod.rs +++ b/src/harness/embeddings/mod.rs @@ -278,8 +278,8 @@ mod cohere; mod noop; mod ollama; mod openai; -pub mod rate_limit; -pub mod retry_after; +mod rate_limit; +mod retry_after; mod voyage; pub use noop::NoopEmbeddingModel; @@ -287,6 +287,10 @@ pub use ollama::{ DEFAULT_OLLAMA_DIMENSIONS, DEFAULT_OLLAMA_MODEL, DEFAULT_OLLAMA_URL, OllamaEmbeddingModel, }; pub use openai::OpenAiEmbeddingModel; +pub use rate_limit::{DEFAULT_REQUESTS_PER_MINUTE, acquire, rate_limit, set_rate_limit}; +pub use retry_after::{ + BASE_BACKOFF_MS, MAX_BACKOFF_MS, MAX_RETRIES, backoff_ms_for_attempt, parse_retry_after_ms, +}; pub use types::format_embedding_signature; pub use voyage::{ VOYAGE_API_BASE, VOYAGE_DEFAULT_DIMENSIONS, VOYAGE_DEFAULT_MODEL, VoyageEmbeddingModel, diff --git a/src/harness/embeddings/noop.rs b/src/harness/embeddings/noop.rs index 2133319..b2f4c4d 100644 --- a/src/harness/embeddings/noop.rs +++ b/src/harness/embeddings/noop.rs @@ -23,7 +23,7 @@ impl EmbeddingModel for NoopEmbeddingModel { 0 } - async fn embed(&self, _texts: &[String]) -> Result>> { - Ok(Vec::new()) + async fn embed(&self, texts: &[String]) -> Result>> { + Ok(vec![Vec::new(); texts.len()]) } } diff --git a/src/harness/embeddings/openai.rs b/src/harness/embeddings/openai.rs index 1fab0da..71c92d4 100644 --- a/src/harness/embeddings/openai.rs +++ b/src/harness/embeddings/openai.rs @@ -65,7 +65,7 @@ impl OpenAiEmbeddingModel { /// Overrides the embedding model id. pub fn with_model(mut self, model: impl Into) -> Self { - self.model = normalize_model_for_base_url(&self.base_url, &model.into()); + self.model = model.into(); self } @@ -73,7 +73,11 @@ impl OpenAiEmbeddingModel { /// endpoint is always `{base_url}/embeddings`. pub fn with_base_url(mut self, base_url: impl Into) -> Self { self.base_url = base_url.into().trim_end_matches('/').to_string(); - self.model = normalize_model_for_base_url(&self.base_url, &self.model); + self + } + + pub(crate) fn with_client(mut self, client: reqwest::Client) -> Self { + self.client = client; self } @@ -287,22 +291,3 @@ fn embeddings_url(base_url: &str) -> String { format!("{}/embeddings", base_url.trim_end_matches('/')) } } - -fn normalize_model_for_base_url(base_url: &str, model: &str) -> String { - let is_gemini = reqwest::Url::parse(base_url) - .ok() - .and_then(|url| url.host_str().map(str::to_ascii_lowercase)) - .is_some_and(|host| { - host == "generativelanguage.googleapis.com" - || host.ends_with(".generativelanguage.googleapis.com") - }); - if is_gemini - && !model.is_empty() - && !model.starts_with("models/") - && !model.starts_with("tunedModels/") - { - format!("models/{model}") - } else { - model.to_owned() - } -} diff --git a/src/harness/embeddings/rate_limit.rs b/src/harness/embeddings/rate_limit.rs index 4e90dfd..4d02772 100644 --- a/src/harness/embeddings/rate_limit.rs +++ b/src/harness/embeddings/rate_limit.rs @@ -15,13 +15,13 @@ static BUCKETS: OnceLock>>> = OnceLock::n pub fn set_rate_limit(per_minute: u32) { let previous = CONFIGURED_LIMIT.swap(per_minute, Ordering::Relaxed); - if previous != per_minute { - if let Some(registry) = BUCKETS.get() { - registry - .lock() - .unwrap_or_else(PoisonError::into_inner) - .clear(); - } + if previous != per_minute + && let Some(registry) = BUCKETS.get() + { + registry + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clear(); } tracing::debug!( target: "tinyagents::embeddings::rate_limit", diff --git a/src/harness/embeddings/test.rs b/src/harness/embeddings/test.rs index f6555bb..42559ec 100644 --- a/src/harness/embeddings/test.rs +++ b/src/harness/embeddings/test.rs @@ -270,8 +270,8 @@ fn embedding_identity_signature_is_stable() { ); } -#[test] -fn voyage_and_noop_identity_match_host_contract() { +#[tokio::test] +async fn voyage_and_noop_identity_match_host_contract() { let voyage = VoyageEmbeddingModel::new("test-key"); assert_eq!(voyage.name(), "voyage"); assert_eq!(voyage.model_id(), VOYAGE_DEFAULT_MODEL); @@ -282,4 +282,18 @@ fn voyage_and_noop_identity_match_host_contract() { let noop = NoopEmbeddingModel; assert_eq!(noop.signature(), "provider=none;model=none;dims=0"); + assert_eq!( + noop.embed(&["first".into(), "second".into()]) + .await + .unwrap(), + vec![Vec::::new(), Vec::::new()] + ); +} + +#[test] +fn gemini_openai_compatible_model_id_is_not_rewritten() { + let model = OpenAiEmbeddingModel::new("test-key") + .with_base_url("https://generativelanguage.googleapis.com/v1beta/openai") + .with_model("gemini-embedding-001"); + assert_eq!(model.model(), "gemini-embedding-001"); } From 777cd7600c4b6560ec6e9732b4f06d462c6059ed Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sun, 12 Jul 2026 11:59:01 +0000 Subject: [PATCH 3/4] fix(embeddings): preserve asymmetric provider semantics --- src/harness/embeddings/cohere.rs | 18 +++++++++++++++++- src/harness/embeddings/mod.rs | 3 +-- src/harness/embeddings/ollama.rs | 18 ++++++++++-------- src/harness/embeddings/types.rs | 7 +++++++ 4 files changed, 35 insertions(+), 11 deletions(-) diff --git a/src/harness/embeddings/cohere.rs b/src/harness/embeddings/cohere.rs index 4aa21b3..6523ac7 100644 --- a/src/harness/embeddings/cohere.rs +++ b/src/harness/embeddings/cohere.rs @@ -17,6 +17,7 @@ pub struct CohereEmbeddingModel { model: String, dimensions: usize, base_url: String, + query_mode: bool, } impl CohereEmbeddingModel { @@ -27,6 +28,7 @@ impl CohereEmbeddingModel { model: COHERE_DEFAULT_MODEL.to_owned(), dimensions: COHERE_DEFAULT_DIMENSIONS, base_url: COHERE_API_BASE.to_owned(), + query_mode: false, } } @@ -95,8 +97,9 @@ impl EmbeddingModel for CohereEmbeddingModel { let body = serde_json::json!({ "model": self.model, "texts": texts, - "input_type": "search_document", + "input_type": if self.query_mode { "search_query" } else { "search_document" }, "embedding_types": ["float"], + "output_dimension": self.dimensions, }); let mut response = None; @@ -172,6 +175,19 @@ impl EmbeddingModel for CohereEmbeddingModel { } Ok(vectors) } + + async fn embed_query(&self, query: &str) -> Result> { + let query_model = Self { + client: self.client.clone(), + api_key: self.api_key.clone(), + model: self.model.clone(), + dimensions: self.dimensions, + base_url: self.base_url.clone(), + query_mode: true, + }; + let mut vectors = query_model.embed(&[query.to_owned()]).await?; + Ok(vectors.pop().unwrap_or_default()) + } } #[cfg(test)] diff --git a/src/harness/embeddings/mod.rs b/src/harness/embeddings/mod.rs index 2b54448..13c294f 100644 --- a/src/harness/embeddings/mod.rs +++ b/src/harness/embeddings/mod.rs @@ -267,8 +267,7 @@ impl Retriever { /// with a different embedding model. An empty store never errors: it /// answers every query with no hits. pub async fn retrieve(&self, query: &str, top_k: usize) -> Result> { - let mut vectors = self.model.embed(&[query.to_string()]).await?; - let query_vector = vectors.pop().unwrap_or_default(); + let query_vector = self.model.embed_query(query).await?; self.store.query(&query_vector, top_k).await } } diff --git a/src/harness/embeddings/ollama.rs b/src/harness/embeddings/ollama.rs index 96ad6cb..2b6ed3b 100644 --- a/src/harness/embeddings/ollama.rs +++ b/src/harness/embeddings/ollama.rs @@ -70,7 +70,7 @@ impl OllamaEmbeddingModel { }) } - async fn embed_one_with_nan_recovery(&self, text: &str) -> Result>> { + async fn embed_one_with_nan_recovery(&self, text: &str) -> Result> { let response = self.request(vec![text.to_owned()]).await?; if !response.status().is_success() { let status = response.status(); @@ -79,9 +79,11 @@ impl OllamaEmbeddingModel { tracing::warn!( target: "tinyagents::embeddings::ollama", model = self.model, - "[embeddings] substituting empty vector for Ollama NaN response" + "[embeddings] Ollama input produced NaN" ); - return Ok(None); + return Err(TinyAgentsError::Embedding( + "Ollama could not encode input without NaN values".into(), + )); } return Err(ollama_http_error(status, &body)); } @@ -94,7 +96,7 @@ impl OllamaEmbeddingModel { } let vector = payload.embeddings.into_iter().next().unwrap(); self.validate_dimensions(0, &vector)?; - Ok(Some(vector)) + Ok(vector) } async fn embed_per_text( @@ -104,9 +106,7 @@ impl OllamaEmbeddingModel { ) -> Result>> { let mut output = vec![Vec::new(); total]; for (index, text) in live { - if let Some(vector) = self.embed_one_with_nan_recovery(text).await? { - output[*index] = vector; - } + output[*index] = self.embed_one_with_nan_recovery(text).await?; } Ok(output) } @@ -273,7 +273,9 @@ impl EmbeddingModel for OllamaEmbeddingModel { "[embeddings] recovering Ollama NaN batch per text" ); if live.len() == 1 { - return Ok(vec![Vec::new(); texts.len()]); + return Err(TinyAgentsError::Embedding( + "Ollama could not encode input without NaN values".into(), + )); } return self.embed_per_text(texts.len(), &live).await; } diff --git a/src/harness/embeddings/types.rs b/src/harness/embeddings/types.rs index 6b2b138..d4c573e 100644 --- a/src/harness/embeddings/types.rs +++ b/src/harness/embeddings/types.rs @@ -50,6 +50,13 @@ pub trait EmbeddingModel: Send + Sync { /// Returning an empty `Vec` for empty input is valid. async fn embed(&self, texts: &[String]) -> Result>>; + /// Embeds a retrieval query. Asymmetric providers can override this; + /// symmetric models reuse [`Self::embed`]. + async fn embed_query(&self, query: &str) -> Result> { + let mut vectors = self.embed(&[query.to_owned()]).await?; + Ok(vectors.pop().unwrap_or_default()) + } + /// Returns the fixed dimensionality of every vector this model produces. fn dimensions(&self) -> usize; From 8f753557278e448c6c5bcf48ab8503505a47c663 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sun, 12 Jul 2026 12:07:41 +0000 Subject: [PATCH 4/4] fix(embeddings): validate provider-specific inputs --- src/harness/embeddings/cohere.rs | 6 ++++-- src/harness/embeddings/ollama.rs | 5 +++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/src/harness/embeddings/cohere.rs b/src/harness/embeddings/cohere.rs index 6523ac7..9b36250 100644 --- a/src/harness/embeddings/cohere.rs +++ b/src/harness/embeddings/cohere.rs @@ -94,13 +94,15 @@ impl EmbeddingModel for CohereEmbeddingModel { } let url = format!("{}/v2/embed", self.base_url); - let body = serde_json::json!({ + let mut body = serde_json::json!({ "model": self.model, "texts": texts, "input_type": if self.query_mode { "search_query" } else { "search_document" }, "embedding_types": ["float"], - "output_dimension": self.dimensions, }); + if self.model.to_ascii_lowercase().contains("v4") && self.dimensions > 0 { + body["output_dimension"] = serde_json::json!(self.dimensions); + } let mut response = None; for attempt in 0..=MAX_RETRIES { diff --git a/src/harness/embeddings/ollama.rs b/src/harness/embeddings/ollama.rs index 2b6ed3b..4c45b9a 100644 --- a/src/harness/embeddings/ollama.rs +++ b/src/harness/embeddings/ollama.rs @@ -257,6 +257,11 @@ impl EmbeddingModel for OllamaEmbeddingModel { if live.is_empty() { return Ok(vec![Vec::new(); texts.len()]); } + if live.len() != texts.len() { + return Err(TinyAgentsError::Validation( + "Ollama embedding batches must not mix blank and nonblank inputs".into(), + )); + } let input = live .iter() .map(|(_, text)| text.clone())