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..0c2c664 --- /dev/null +++ b/src/harness/embeddings/cloud.rs @@ -0,0 +1,124 @@ +//! 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 { + client: reqwest::Client, + 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 { + client: reqwest::Client::new(), + 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_client(self.client.clone()) + .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..9b36250 --- /dev/null +++ b/src/harness/embeddings/cohere.rs @@ -0,0 +1,223 @@ +//! 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, + query_mode: bool, +} + +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(), + query_mode: false, + } + } + + 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 mut body = serde_json::json!({ + "model": self.model, + "texts": texts, + "input_type": if self.query_mode { "search_query" } else { "search_document" }, + "embedding_types": ["float"], + }); + 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 { + 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) + } + + 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)] +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..13c294f 100644 --- a/src/harness/embeddings/mod.rs +++ b/src/harness/embeddings/mod.rs @@ -267,15 +267,39 @@ 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 } } +mod cloud; +mod cohere; +mod noop; +mod ollama; mod openai; +mod rate_limit; +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 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, +}; #[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..b2f4c4d --- /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![Vec::new(); texts.len()]) + } +} diff --git a/src/harness/embeddings/ollama.rs b/src/harness/embeddings/ollama.rs new file mode 100644 index 0000000..4c45b9a --- /dev/null +++ b/src/harness/embeddings/ollama.rs @@ -0,0 +1,341 @@ +//! 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] Ollama input produced NaN" + ); + return Err(TinyAgentsError::Embedding( + "Ollama could not encode input without NaN values".into(), + )); + } + 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(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 { + output[*index] = self.embed_one_with_nan_recovery(text).await?; + } + 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()]); + } + 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()) + .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 Err(TinyAgentsError::Embedding( + "Ollama could not encode input without NaN values".into(), + )); + } + 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..71c92d4 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,6 +58,8 @@ impl OpenAiEmbeddingModel { model: DEFAULT_MODEL.to_string(), base_url: DEFAULT_BASE_URL.to_string(), dimensions: DEFAULT_DIMENSIONS, + send_dimensions: true, + requires_api_key: true, } } @@ -71,6 +76,11 @@ impl OpenAiEmbeddingModel { self } + pub(crate) fn with_client(mut self, client: reqwest::Client) -> Self { + self.client = client; + self + } + /// Overrides the reported dimensionality (and requests it from the API /// via the `dimensions` parameter, which `text-embedding-3-*` supports). pub fn with_dimensions(mut self, dimensions: usize) -> Self { @@ -78,6 +88,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 +146,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 +244,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 +277,17 @@ 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('/')) + } +} diff --git a/src/harness/embeddings/rate_limit.rs b/src/harness/embeddings/rate_limit.rs new file mode 100644 index 0000000..4d02772 --- /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 + && 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..42559ec 100644 --- a/src/harness/embeddings/test.rs +++ b/src/harness/embeddings/test.rs @@ -255,3 +255,45 @@ 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" + ); +} + +#[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); + 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"); + 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"); +} diff --git a/src/harness/embeddings/types.rs b/src/harness/embeddings/types.rs index 2c38d03..d4c573e 100644 --- a/src/harness/embeddings/types.rs +++ b/src/harness/embeddings/types.rs @@ -39,13 +39,38 @@ 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. 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; + + /// 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 +140,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 + } +}