diff --git a/Cargo.lock b/Cargo.lock index d9a7fec9..240bf5fc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1285,6 +1285,8 @@ name = "tinyinference-embeddings" version = "0.3.0" dependencies = [ "async-trait", + "futures", + "http", "once_cell", "regex", "reqwest", diff --git a/README.md b/README.md index 5879e391..e23fc8ef 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,8 @@ The workspace provides: and Mistral; - OpenAI, Cohere, Ollama, Voyage, cloud, no-op, and deterministic mock embeddings; +- standalone document reranking through Voyage, with original-position results, + cancellation, bounded retries, and provider-reported usage; - request caching, stream accumulation, normalized provider failures, provider-neutral retry classification and `Retry-After` parsing; - conservative context-window and vision-capability hints for raw model ids @@ -52,7 +54,7 @@ assert_eq!(response.text(), "hello"); TinyAgents vendors this repository at `vendor/tinyinference` and re-exports the public modules through its historical `tinyagents::harness::*` paths. New code can depend on `tinyinference-llm` for language models, -`tinyinference-embeddings` for vector generation and retrieval, +`tinyinference-embeddings` for vector generation, retrieval, and reranking, `tinyinference-local` for endpoint-only local inference, `tinyinference-providers` for provider authentication and routing primitives, `tinyinference-voice` for speech inference and streaming-audio mechanics, @@ -62,6 +64,54 @@ standards and OpenRouter media transport, download, resume), `tinyinference-decisions` for typed Jev and Sage decisions, and `tinyinference-core` only for shared infrastructure. +### Voyage reranking + +`tinyinference-embeddings` exports `Reranker`, `RerankRequest`, +`VoyageReranker`, and `VoyageRerankConfig`. Supply your Voyage API key explicitly: + +```rust,no_run +use tinyinference_embeddings::{Reranker, RerankRequest, VoyageReranker}; + +# async fn example(api_key: String) -> tinyinference_embeddings::Result<()> { +let model = VoyageReranker::new(api_key)?; +let documents = vec!["Reset your password in Settings.".into(), "Shipping takes two days.".into()]; +let mut request = RerankRequest::new("How do I reset my password?", documents); +request.top_k = Some(1); +let response = model.rerank(request).await?; +for result in response.results { + println!("original position: {}, score: {}", result.index, result.relevance_score); +} +# Ok(()) +# } +``` + +Use candidate texts from any search system. Our `Retriever` does not preserve +source text: resolve its returned IDs against your own document storage first. +Missing text should fail before the request rather than silently removing a +candidate. Returned indices address the exact submitted list. Keep reranking +scores separate from `ScoredDoc.score`, which remains vector similarity. +The `rerank` module's compiled rustdoc shows this mapping and caller-owned fallback. + +Defaults are `rerank-3`, a 30-second total deadline, up to three retries of +429/500/502/503/504 responses, 8 MiB serialized requests, and 2 MiB responses. +`VoyageRerankConfig` makes these configurable. Requests accept up to 1,000 +documents; zero results or empty candidates require no provider call. +Blank queries/documents are rejected for nonempty requests. Provider token +limits still apply. Text shortening is disabled unless `request.truncate` is +explicitly enabled; large batches are never silently split or dropped. + +Clone `request.cancellation` to cancel while waiting for pacing, a response, +or a retry. The deadline covers all attempts. Transport interruptions are not +replayed; the provider may still finish and bill an interrupted request. +Explicit status retries can repeat billable work; set `max_retries` to zero to +disable them. Missing usage is unknown, and reported usage covers only the +successful response. Decide in your app whether provider failure preserves +original search order; propagate cancellation rather than treating it as fallback. + +Errors use `Error::Rerank(RerankError)` for provider, timeout, and response +failures; `Validation` and `Cancelled` retain their existing roles. Downstream +exhaustive matches on the crate's `Error` must add the new `Rerank` arm. + ### Levanto Sage Sage uses typed decisions. A missing verdict means Sage is unsure and should @@ -114,7 +164,7 @@ crates/tinyinference-llm/ ├── tool.rs model-visible tool schemas and call/delta shapes └── usage/ normalized token accounting crates/tinyinference-embeddings/ -└── src/ embedding clients, vector store, and retriever +└── src/ embedding clients, rerank/, vector store, and retriever crates/tinyinference-local/ └── src/ local endpoint probing, model selection, and inference crates/tinyinference-providers/ diff --git a/crates/tinyinference-embeddings/Cargo.toml b/crates/tinyinference-embeddings/Cargo.toml index e9288a95..87b5240f 100644 --- a/crates/tinyinference-embeddings/Cargo.toml +++ b/crates/tinyinference-embeddings/Cargo.toml @@ -24,5 +24,10 @@ tokio = { workspace = true } tracing = { workspace = true } url = { workspace = true } +[dev-dependencies] +futures = { workspace = true } +http = "1" +tokio = { workspace = true, features = ["test-util"] } + [lints] workspace = true diff --git a/crates/tinyinference-embeddings/src/error.rs b/crates/tinyinference-embeddings/src/error.rs index 8a8de872..055cf901 100644 --- a/crates/tinyinference-embeddings/src/error.rs +++ b/crates/tinyinference-embeddings/src/error.rs @@ -20,4 +20,7 @@ pub enum Error { /// The caller cancelled an embedding request. #[error("embedding request cancelled")] Cancelled, + /// Reranking failed at the provider boundary. + #[error(transparent)] + Rerank(#[from] crate::rerank::RerankError), } diff --git a/crates/tinyinference-embeddings/src/lib.rs b/crates/tinyinference-embeddings/src/lib.rs index 89a75c6a..10767936 100644 --- a/crates/tinyinference-embeddings/src/lib.rs +++ b/crates/tinyinference-embeddings/src/lib.rs @@ -55,10 +55,15 @@ pub mod catalog; mod error; pub mod factory; pub mod probe; +pub mod rerank; pub mod served_models; mod types; pub use error::{Error, Result}; +pub use rerank::{ + RerankCancellation, RerankError, RerankRequest, RerankResponse, RerankResult, RerankUsage, + Reranker, VoyageRerankConfig, VoyageReranker, +}; pub use types::*; use async_trait::async_trait; diff --git a/crates/tinyinference-embeddings/src/rerank/mod.rs b/crates/tinyinference-embeddings/src/rerank/mod.rs new file mode 100644 index 00000000..d0662f93 --- /dev/null +++ b/crates/tinyinference-embeddings/src/rerank/mod.rs @@ -0,0 +1,131 @@ +//! Provider-neutral reranking of candidate document texts. +//! +//! [`VoyageReranker`] accepts an explicit API key and defaults to `rerank-3`. +//! Candidates may come from any search system. Results address the original +//! text list; keep document IDs, metadata, and vector similarity scores with +//! the caller. Truncation is disabled unless explicitly requested. +//! +//! A total deadline includes pacing, HTTP requests, response reads, and retry +//! waits. Transport interruptions are not retried because the provider may +//! already have billed the request. Explicit temporary HTTP statuses may be +//! retried; configure zero retries to disable this. Reported usage covers only +//! the successful response and missing usage means unknown, never zero. +//! +//! # Rerank search hits using caller-owned source text +//! +//! [`crate::Retriever`] stores vectors and metadata, not document text. Resolve +//! texts by ID before sending any request; missing text should not silently +//! remove a candidate and change the index mapping. +//! +//! ```no_run +//! use std::collections::HashMap; +//! use tinyinference_embeddings::{ +//! Error, Result, Retriever, Reranker, RerankRequest, ScoredDoc, +//! }; +//! +//! async fn search( +//! retriever: &Retriever, +//! reranker: &dyn Reranker, +//! texts: &HashMap, +//! query: &str, +//! ) -> Result)>> { +//! let hits = retriever.retrieve(query, 30).await?; +//! let documents = hits.iter().map(|hit| { +//! texts.get(&hit.id).cloned() +//! .ok_or_else(|| Error::Validation("missing source document".into())) +//! }).collect::>>()?; +//! let mut request = RerankRequest::new(query, documents); +//! request.top_k = Some(5); +//! match reranker.rerank(request).await { +//! Ok(response) => response.results.into_iter().map(|result| { +//! let hit = hits.get(result.index).cloned() +//! .ok_or_else(|| Error::Validation("invalid reranking index".into()))?; +//! Ok((hit, Some(result.relevance_score))) +//! }).collect(), +//! Err(Error::Cancelled) => Err(Error::Cancelled), +//! // This caller elects to retain search order on provider failure. +//! Err(Error::Rerank(_)) => Ok(hits.into_iter().take(5) +//! .map(|hit| (hit, None)).collect()), +//! Err(error) => Err(error), +//! } +//! } +//! ``` + +mod types; +mod voyage; + +pub use types::{ + RerankCancellation, RerankError, RerankRequest, RerankResponse, RerankResult, RerankUsage, + Reranker, VoyageRerankConfig, +}; +pub use voyage::VoyageReranker; + +fn validate_request(request: &RerankRequest) -> crate::Result { + let count = request.documents.len(); + let top_k = request.top_k.unwrap_or(count).min(count); + if top_k == 0 { + return Ok(0); + } + if request.query.trim().is_empty() { + return Err(crate::Error::Validation( + "reranking query must not be blank".into(), + )); + } + if let Some(index) = request + .documents + .iter() + .position(|text| text.trim().is_empty()) + { + return Err(crate::Error::Validation(format!( + "reranking document at index {index} must not be blank" + ))); + } + Ok(top_k) +} + +fn validate_results( + response: &mut RerankResponse, + document_count: usize, + top_k: usize, +) -> crate::Result<()> { + if response.results.len() != top_k { + return Err(RerankError::InvalidResponse { + reason: "unexpected result count", + } + .into()); + } + let mut seen = vec![false; document_count]; + for result in &response.results { + let Some(already_seen) = seen.get_mut(result.index) else { + return Err(RerankError::InvalidResponse { + reason: "document index out of range", + } + .into()); + }; + if *already_seen { + return Err(RerankError::InvalidResponse { + reason: "duplicate document index", + } + .into()); + } + *already_seen = true; + if !result.relevance_score.is_finite() { + return Err(RerankError::InvalidResponse { + reason: "non-finite relevance score", + } + .into()); + } + } + response.results.sort_by(|a, b| { + if a.relevance_score == b.relevance_score { + a.index.cmp(&b.index) + } else { + b.relevance_score.total_cmp(&a.relevance_score) + } + }); + Ok(()) +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyinference-embeddings/src/rerank/mod_tests.rs b/crates/tinyinference-embeddings/src/rerank/mod_tests.rs new file mode 100644 index 00000000..7680b73a --- /dev/null +++ b/crates/tinyinference-embeddings/src/rerank/mod_tests.rs @@ -0,0 +1,196 @@ +use super::*; + +use crate::{Error, Result, ScoredDoc}; +use serde_json::json; +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, +}; + +#[tokio::test] +async fn empty_candidates_need_no_query_or_provider_call() { + let reranker = VoyageReranker::new("test-key").unwrap(); + let response = reranker + .rerank(RerankRequest::new("", Vec::new())) + .await + .unwrap(); + assert_eq!( + response, + RerankResponse { + results: vec![], + usage: None + } + ); +} + +#[test] +fn results_have_a_stable_serializable_shape() { + let response = RerankResponse { + results: vec![RerankResult { + index: 2, + relevance_score: 0.75, + }], + usage: Some(RerankUsage { input_tokens: 17 }), + }; + assert_eq!( + serde_json::to_value(response).unwrap(), + json!({ + "results":[{"index":2,"relevance_score":0.75}], "usage":{"input_tokens":17} + }) + ); + fn send_sync() {} + send_sync::(); + send_sync::(); + send_sync::>(); +} + +struct FixedReranker { + documents: Vec, + reply: Mutex>>, +} + +#[async_trait::async_trait] +impl Reranker for FixedReranker { + fn name(&self) -> &str { + "fixture" + } + fn model_id(&self) -> &str { + "fixed" + } + async fn rerank(&self, request: RerankRequest) -> Result { + assert_eq!(request.documents, self.documents); + self.reply.lock().unwrap().take().expect("only one call") + } +} + +async fn rerank_hits( + model: &dyn Reranker, + hits: Vec, + texts: &HashMap, +) -> Result)>> { + let documents = hits + .iter() + .map(|hit| { + texts + .get(&hit.id) + .cloned() + .ok_or_else(|| Error::Validation("missing source document".into())) + }) + .collect::>>()?; + match model.rerank(RerankRequest::new("query", documents)).await { + Ok(response) => response + .results + .into_iter() + .map(|result| { + let hit = hits + .get(result.index) + .cloned() + .ok_or_else(|| Error::Validation("invalid result index".into()))?; + Ok((hit, Some(result.relevance_score))) + }) + .collect(), + Err(Error::Cancelled) => Err(Error::Cancelled), + Err(Error::Rerank(_)) => Ok(hits.into_iter().map(|hit| (hit, None)).collect()), + Err(error) => Err(error), + } +} + +fn reverse_rankings(documents: Vec) -> FixedReranker { + FixedReranker { + documents, + reply: Mutex::new(Some(Ok(RerankResponse { + results: vec![ + RerankResult { + index: 1, + relevance_score: 0.9, + }, + RerankResult { + index: 0, + relevance_score: 0.1, + }, + ], + usage: None, + }))), + } +} + +#[tokio::test] +async fn external_search_hits_keep_ids_metadata_and_original_scores() { + let hits = vec![ + ScoredDoc { + id: "a".into(), + score: 0.8, + metadata: json!({"source":"one"}), + }, + ScoredDoc { + id: "b".into(), + score: 0.6, + metadata: json!({"source":"two"}), + }, + ]; + let texts = HashMap::from([("a".into(), "first".into()), ("b".into(), "second".into())]); + let model = reverse_rankings(vec!["first".into(), "second".into()]); + let result = rerank_hits(&model, hits.clone(), &texts).await.unwrap(); + assert_eq!( + result, + vec![(hits[1].clone(), Some(0.9)), (hits[0].clone(), Some(0.1))] + ); +} + +#[tokio::test] +async fn retriever_hits_use_explicit_source_text_and_missing_text_prevents_reranking() { + let retriever = crate::Retriever::new( + Arc::new(crate::MockEmbeddingModel::new(16)), + Arc::new(crate::InMemoryVectorStore::new()), + ); + retriever + .index(vec![ + ("a".into(), "alpha".into(), json!({"source":"one"})), + ("b".into(), "beta".into(), json!({"source":"two"})), + ]) + .await + .unwrap(); + let hits = retriever.retrieve("alpha", 2).await.unwrap(); + assert_eq!(hits[0].id, "a"); + let texts = HashMap::from([("a".into(), "alpha".into()), ("b".into(), "beta".into())]); + let model = reverse_rankings(vec!["alpha".into(), "beta".into()]); + let result = rerank_hits(&model, hits.clone(), &texts).await.unwrap(); + assert_eq!( + result, + vec![(hits[1].clone(), Some(0.9)), (hits[0].clone(), Some(0.1))] + ); + let unused = reverse_rankings(vec![]); + assert!(matches!( + rerank_hits(&unused, hits, &HashMap::new()).await, + Err(Error::Validation(_)) + )); + assert!(unused.reply.lock().unwrap().is_some()); +} + +#[tokio::test] +async fn caller_can_fall_back_on_provider_failure_but_preserves_cancellation() { + let hit = ScoredDoc { + id: "a".into(), + score: 0.8, + metadata: json!({"source":"one"}), + }; + let texts = HashMap::from([("a".into(), "first".into())]); + let failed = FixedReranker { + documents: vec!["first".into()], + reply: Mutex::new(Some(Err(RerankError::Transport.into()))), + }; + assert_eq!( + rerank_hits(&failed, vec![hit.clone()], &texts) + .await + .unwrap(), + vec![(hit.clone(), None)] + ); + let cancelled = FixedReranker { + documents: vec!["first".into()], + reply: Mutex::new(Some(Err(Error::Cancelled))), + }; + assert!(matches!( + rerank_hits(&cancelled, vec![hit], &texts).await, + Err(Error::Cancelled) + )); +} diff --git a/crates/tinyinference-embeddings/src/rerank/types.rs b/crates/tinyinference-embeddings/src/rerank/types.rs new file mode 100644 index 00000000..10c0fdb7 --- /dev/null +++ b/crates/tinyinference-embeddings/src/rerank/types.rs @@ -0,0 +1,187 @@ +//! Public reranking requests, results, and configuration. + +use std::{fmt, time::Duration}; + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; + +/// Cooperative cancellation shared with the caller; cancelling stops local +/// waiting, but cannot guarantee cancellation of provider processing or billing. +pub type RerankCancellation = crate::EmbeddingCancellation; + +/// Candidate texts in the exact order addressed by returned result indices. +/// +/// Keep IDs, metadata, and original search scores with the caller. Repeated +/// document text at different positions represents distinct candidates. +#[derive(Clone)] +pub struct RerankRequest { + /// Question or search query used to score the documents. + pub query: String, + /// Original candidate text; empty or whitespace-only entries are invalid. + pub documents: Vec, + /// Maximum results; `None` returns all, zero returns none, and values above + /// the candidate count are clamped to that count. + pub top_k: Option, + /// Explicitly permit provider-side shortening of oversized input. + pub truncate: bool, + /// Caller-controlled cancellation signal. + pub cancellation: RerankCancellation, +} + +impl RerankRequest { + /// Creates a request for all candidates with truncation disabled. + pub fn new(query: impl Into, documents: Vec) -> Self { + Self { + query: query.into(), + documents, + top_k: None, + truncate: false, + cancellation: RerankCancellation::new(), + } + } +} + +impl fmt::Debug for RerankRequest { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RerankRequest") + .field("documents", &self.documents.len()) + .field("top_k", &self.top_k) + .field("truncate", &self.truncate) + .field("cancelled", &self.cancellation.is_cancelled()) + .finish_non_exhaustive() + } +} + +/// One result addressing the caller's original candidate list. +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] +pub struct RerankResult { + /// Zero-based position in the submitted documents, not the sorted results. + pub index: usize, + /// Finite provider score; higher is better. This is neither a probability + /// nor the vector similarity stored in [`crate::ScoredDoc::score`]. + pub relevance_score: f64, +} + +/// Measured usage from the successful provider response only. +/// +/// Failed or retried attempts may incur additional unreported usage. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct RerankUsage { + /// Provider-reported input tokens, without local estimates. + pub input_tokens: u64, +} + +/// Validated results in descending score order, ties ordered by original index. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct RerankResponse { + /// Exactly the effective requested result count; never a partial success. + pub results: Vec, + /// Missing or invalid provider usage is unknown, not zero. + pub usage: Option, +} + +/// A provider-neutral reranking model, usable behind an `Arc`. +#[async_trait] +pub trait Reranker: Send + Sync { + /// Stable provider name. + fn name(&self) -> &str; + /// Configured provider model identifier. + fn model_id(&self) -> &str; + /// Reranks original candidate texts without modifying caller-owned metadata. + /// + /// Implementations must preserve original indices, return finite scores and + /// exactly the effective result count, and respect cancellation. Empty + /// candidates or zero results perform no provider call; pre-cancellation + /// still wins. Nonempty requests require nonblank query and documents. + /// + /// # Errors + /// Returns [`crate::Error::Validation`] for invalid inputs, + /// [`crate::Error::Cancelled`] for cancellation, or + /// [`crate::Error::Rerank`] for provider, deadline, or response failures. + /// Fallback to original search order belongs to the caller. + async fn rerank(&self, request: RerankRequest) -> crate::Result; +} + +/// Voyage settings, validated when constructing [`super::VoyageReranker`]. +#[derive(Clone)] +pub struct VoyageRerankConfig { + /// Model ID; defaults to `rerank-3`. Nonblank custom IDs are accepted. + pub model: String, + /// API base URL without `/rerank`; defaults to Voyage's `/v1` URL. + pub base_url: String, + /// Total budget including request pacing, retries, and response reads. + pub timeout: Duration, + /// Extra attempts for 429/500/502/503/504, from zero to the core retry cap. + /// Retries may repeat billable work; zero disables them. + pub max_retries: u32, + /// Maximum serialized request size, including JSON escaping. + pub max_request_bytes: usize, + /// Maximum response body size, enforced while reading. + pub max_response_bytes: usize, +} + +impl Default for VoyageRerankConfig { + fn default() -> Self { + Self { + model: "rerank-3".into(), + base_url: crate::VOYAGE_API_BASE.into(), + timeout: Duration::from_secs(30), + max_retries: tinyinference_core::MAX_RETRIES, + max_request_bytes: 8 * 1024 * 1024, + max_response_bytes: 2 * 1024 * 1024, + } + } +} + +impl fmt::Debug for VoyageRerankConfig { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("VoyageRerankConfig") + .field("timeout", &self.timeout) + .field("max_retries", &self.max_retries) + .field("max_request_bytes", &self.max_request_bytes) + .field("max_response_bytes", &self.max_response_bytes) + .finish_non_exhaustive() + } +} + +/// Distinguishable failures of a reranking operation. +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum RerankError { + /// Credentials or access were rejected. + #[error("reranking authentication rejected (HTTP {status})")] + Authentication { + /// Original 401 or 403 status. + status: u16, + }, + /// The provider still rate-limited the request after allowed retries. + #[error("reranking rate limit exceeded")] + RateLimited { + /// Provider delay in milliseconds, bounded by the core retry policy. + retry_after_ms: Option, + }, + /// Another unsuccessful provider status. + #[error("reranking provider returned HTTP {status}")] + HttpStatus { + /// Original HTTP status. + status: u16, + }, + /// The total operation deadline expired. + #[error("reranking request timed out")] + Timeout, + /// The request or response was interrupted; provider billing may be unknown. + #[error("reranking transport failed")] + Transport, + /// The response could not satisfy the requested result contract. + #[error("invalid reranking response: {reason}")] + InvalidResponse { + /// Credential-safe explanation without provider body or document text. + reason: &'static str, + }, + /// The response exceeded the configured body limit. + #[error("reranking response exceeded {limit} bytes")] + ResponseTooLarge { + /// Configured maximum response bytes. + limit: usize, + }, +} diff --git a/crates/tinyinference-embeddings/src/rerank/voyage.rs b/crates/tinyinference-embeddings/src/rerank/voyage.rs new file mode 100644 index 00000000..edc62acc --- /dev/null +++ b/crates/tinyinference-embeddings/src/rerank/voyage.rs @@ -0,0 +1,328 @@ +use std::{fmt, io::Write, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use reqwest::header::{AUTHORIZATION, CONTENT_TYPE, HeaderValue}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use tokio::time::Instant; + +use super::{ + RerankCancellation, RerankError, RerankRequest, RerankResponse, RerankResult, RerankUsage, + Reranker, VoyageRerankConfig, +}; +use crate::{Error, Result}; + +/// A Voyage reranker with caller-supplied credentials and bounded operations. +/// +/// The client never reads ambient credentials. Timeout or cancellation stops +/// local waiting; Voyage may still finish and bill work already submitted. +/// Automatic retries cover explicit 429/500/502/503/504 responses only, and may +/// repeat billed work. Configure zero retries to disable them. +pub struct VoyageReranker { + config: VoyageRerankConfig, + authorization: HeaderValue, + client: reqwest::Client, + transport: Arc, +} + +impl fmt::Debug for VoyageReranker { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("VoyageReranker") + .field("config", &self.config) + .finish_non_exhaustive() + } +} + +impl VoyageReranker { + /// Creates a reranker with default settings. + /// + /// # Errors + /// Returns [`Error::Validation`] for invalid credentials or + /// [`Error::Rerank`] if the HTTP client cannot be created. + pub fn new(api_key: impl Into) -> Result { + Self::with_options(api_key, VoyageRerankConfig::default()) + } + + /// Creates a reranker with validated settings and a reusable HTTP client. + /// + /// Credentialed endpoints require HTTPS except on loopback. Userinfo, + /// queries, fragments, and redirects are disallowed. + /// + /// # Errors + /// Returns [`Error::Validation`] for invalid credentials, endpoint, model, + /// timeout, limits, or retry count; [`Error::Rerank`] for client setup failure. + pub fn with_options( + api_key: impl Into, + mut config: VoyageRerankConfig, + ) -> Result { + let api_key = api_key.into(); + if api_key.trim().is_empty() { + return Err(Error::Validation("Voyage API key must not be blank".into())); + } + if config.model.trim().is_empty() + || config.timeout.is_zero() + || config.max_request_bytes == 0 + || config.max_response_bytes == 0 + || config.max_retries > tinyinference_core::MAX_RETRIES + || Instant::now().checked_add(config.timeout).is_none() + { + return Err(Error::Validation( + "invalid reranking model, timeout, byte limit, or retry count".into(), + )); + } + let endpoint = crate::factory::validate_custom_endpoint(&config.base_url, true)?; + let parsed = reqwest::Url::parse(&endpoint) + .map_err(|_| Error::Validation("invalid reranking endpoint".into()))?; + if !parsed.username().is_empty() + || parsed.password().is_some() + || parsed.query().is_some() + || parsed.fragment().is_some() + { + return Err(Error::Validation( + "reranking endpoint must not contain userinfo, query, or fragment".into(), + )); + } + config.base_url = parsed.as_str().trim_end_matches('/').to_owned(); + let mut authorization = HeaderValue::from_str(&format!("Bearer {api_key}")) + .map_err(|_| Error::Validation("invalid Voyage API key header".into()))?; + authorization.set_sensitive(true); + let client = reqwest::Client::builder() + .connect_timeout(config.timeout.min(Duration::from_secs(10))) + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|_| RerankError::Transport)?; + Ok(Self { + config, + authorization, + transport: Arc::new(ReqwestTransport { + client: client.clone(), + }), + client, + }) + } + + async fn call( + &self, + request: &RerankRequest, + top_k: usize, + deadline: Instant, + ) -> Result { + let wire = WireRequest { + query: &request.query, + documents: &request.documents, + model: &self.config.model, + top_k, + return_documents: false, + truncation: request.truncate, + }; + let mut body = BoundedBody { + bytes: Vec::new(), + limit: self.config.max_request_bytes, + }; + serde_json::to_writer(&mut body, &wire).map_err(|_| { + Error::Validation("reranking request exceeds serialized byte limit".into()) + })?; + let outgoing = self + .client + .post(format!("{}/rerank", self.config.base_url)) + .header(AUTHORIZATION, self.authorization.clone()) + .header(CONTENT_TYPE, "application/json") + .body(body.bytes) + .build() + .map_err(|_| Error::Validation("could not construct reranking request".into()))?; + ensure_active(&request.cancellation, deadline)?; + let response = self + .send_with_retry(outgoing, &request.cancellation, deadline) + .await?; + let bytes = read_body(response, self.config.max_response_bytes).await?; + let wire: WireResponse = + serde_json::from_slice(&bytes).map_err(|_| RerankError::InvalidResponse { + reason: "malformed result payload", + })?; + let usage = wire + .usage + .as_ref() + .and_then(|usage| usage.get("total_tokens")) + .and_then(Value::as_u64) + .map(|input_tokens| RerankUsage { input_tokens }); + let mut response = RerankResponse { + results: wire.data, + usage, + }; + super::validate_results(&mut response, request.documents.len(), top_k)?; + Ok(response) + } + + async fn send_with_retry( + &self, + outgoing: reqwest::Request, + cancellation: &RerankCancellation, + deadline: Instant, + ) -> Result { + let mut retries = 0; + loop { + crate::rate_limit::acquire(&self.config.base_url).await; + ensure_active(cancellation, deadline)?; + let attempt = outgoing.try_clone().ok_or_else(|| { + Error::Validation("reranking request body cannot be replayed".into()) + })?; + let response = self.transport.send(attempt).await?; + if response.status().is_success() { + return Ok(response); + } + let status = response.status().as_u16(); + let retry_after = response + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + drop(response); + if matches!(status, 429 | 500 | 502 | 503 | 504) && retries < self.config.max_retries { + let delay = + tinyinference_core::backoff_ms_for_attempt(retries, retry_after.as_deref()); + retries += 1; + tokio::time::sleep(Duration::from_millis(delay)).await; + continue; + } + return Err(match status { + 401 | 403 => RerankError::Authentication { status }, + 429 => RerankError::RateLimited { + retry_after_ms: tinyinference_core::parse_retry_after_ms( + retry_after.as_deref(), + ), + }, + _ => RerankError::HttpStatus { status }, + } + .into()); + } + } + + async fn execute(&self, request: &RerankRequest, deadline: Instant) -> Result { + if request.documents.is_empty() || request.top_k == Some(0) { + return Ok(RerankResponse { + results: vec![], + usage: None, + }); + } + if request.documents.len() > 1_000 { + return Err(Error::Validation( + "Voyage reranking accepts at most 1000 documents".into(), + )); + } + let top_k = super::validate_request(request)?; + ensure_active(&request.cancellation, deadline)?; + self.call(request, top_k, deadline).await + } +} + +#[async_trait] +impl Reranker for VoyageReranker { + fn name(&self) -> &str { + "voyage" + } + fn model_id(&self) -> &str { + &self.config.model + } + + async fn rerank(&self, request: RerankRequest) -> Result { + if request.cancellation.is_cancelled() { + return Err(Error::Cancelled); + } + let deadline = Instant::now() + .checked_add(self.config.timeout) + .ok_or_else(|| Error::Validation("reranking timeout is too large".into()))?; + let result = tokio::select! { + biased; + () = request.cancellation.cancelled() => return Err(Error::Cancelled), + () = tokio::time::sleep_until(deadline) => return Err(RerankError::Timeout.into()), + result = self.execute(&request, deadline) => result, + }; + ensure_active(&request.cancellation, deadline)?; + result + } +} + +fn ensure_active(cancellation: &RerankCancellation, deadline: Instant) -> Result<()> { + if cancellation.is_cancelled() { + Err(Error::Cancelled) + } else if Instant::now() >= deadline { + Err(RerankError::Timeout.into()) + } else { + Ok(()) + } +} + +#[derive(Serialize)] +struct WireRequest<'a> { + query: &'a str, + documents: &'a [String], + model: &'a str, + top_k: usize, + return_documents: bool, + truncation: bool, +} + +#[derive(Deserialize)] +struct WireResponse { + data: Vec, + #[serde(default)] + usage: Option, +} + +struct BoundedBody { + bytes: Vec, + limit: usize, +} + +impl Write for BoundedBody { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + if bytes.len() > self.limit.saturating_sub(self.bytes.len()) { + return Err(std::io::Error::other("request body exceeds limit")); + } + self.bytes.extend_from_slice(bytes); + Ok(bytes.len()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +#[async_trait] +trait HttpTransport: Send + Sync { + async fn send(&self, request: reqwest::Request) -> Result; +} + +struct ReqwestTransport { + client: reqwest::Client, +} + +#[async_trait] +impl HttpTransport for ReqwestTransport { + async fn send(&self, request: reqwest::Request) -> Result { + self.client + .execute(request) + .await + .map_err(|_| RerankError::Transport.into()) + } +} + +async fn read_body(mut response: reqwest::Response, limit: usize) -> Result> { + if response + .content_length() + .is_some_and(|length| length > limit as u64) + { + return Err(RerankError::ResponseTooLarge { limit }.into()); + } + let mut body = Vec::new(); + while let Some(chunk) = response.chunk().await.map_err(|_| RerankError::Transport)? { + if chunk.len() > limit.saturating_sub(body.len()) { + return Err(RerankError::ResponseTooLarge { limit }.into()); + } + body.extend_from_slice(&chunk); + } + Ok(body) +} + +#[cfg(test)] +#[path = "voyage_tests.rs"] +mod tests; diff --git a/crates/tinyinference-embeddings/src/rerank/voyage_tests.rs b/crates/tinyinference-embeddings/src/rerank/voyage_tests.rs new file mode 100644 index 00000000..d0f3f29d --- /dev/null +++ b/crates/tinyinference-embeddings/src/rerank/voyage_tests.rs @@ -0,0 +1,581 @@ +use super::*; + +use serde_json::json; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::{collections::VecDeque, sync::Mutex}; + +const ONE_RESULT: &str = r#"{"data":[{"index":0,"relevance_score":0.75}]}"#; + +fn request() -> RerankRequest { + RerankRequest::new("query", vec!["document".into()]) +} + +struct ScriptedTransport { + responses: Mutex>>, + requests: Mutex>, +} + +#[async_trait] +impl HttpTransport for ScriptedTransport { + async fn send(&self, request: reqwest::Request) -> Result { + self.requests.lock().unwrap().push(request); + self.responses + .lock() + .unwrap() + .pop_front() + .expect("unexpected provider call") + } +} + +fn response(status: u16, body: &str) -> Result { + Ok(http::Response::builder() + .status(status) + .body(body.to_owned()) + .unwrap() + .into()) +} + +fn scripted(responses: Vec>) -> (VoyageReranker, Arc) { + let transport = Arc::new(ScriptedTransport { + responses: Mutex::new(responses.into()), + requests: Mutex::new(vec![]), + }); + let mut model = VoyageReranker::with_options( + "test-key", + VoyageRerankConfig { + base_url: "http://localhost/v1".into(), + ..VoyageRerankConfig::default() + }, + ) + .unwrap(); + model.transport = transport.clone(); + (model, transport) +} + +#[tokio::test] +async fn voyage_preserves_request_text_and_maps_ranked_positions_with_usage() { + let (model, transport) = scripted(vec![response( + 200, + r#"{ + "data":[{"index":1,"relevance_score":0.9},{"index":0,"relevance_score":0.2}], + "usage":{"total_tokens":17},"extra":"ignored" + }"#, + )]); + let mut request = RerankRequest::new(" Which? ", vec!["line\n\"one\"".into(), "café".into()]); + request.top_k = Some(9); + let result = model.rerank(request).await.unwrap(); + assert_eq!( + result.results, + vec![ + super::super::RerankResult { + index: 1, + relevance_score: 0.9 + }, + super::super::RerankResult { + index: 0, + relevance_score: 0.2 + }, + ] + ); + assert_eq!( + result.usage, + Some(super::super::RerankUsage { input_tokens: 17 }) + ); + let requests = transport.requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].method(), reqwest::Method::POST); + assert_eq!(requests[0].url().as_str(), "http://localhost/v1/rerank"); + assert_eq!( + requests[0].headers()[reqwest::header::AUTHORIZATION], + "Bearer test-key" + ); + let body: serde_json::Value = + serde_json::from_slice(requests[0].body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!( + body, + json!({ + "query":" Which? ", "documents":["line\n\"one\"", "café"], "model":"rerank-3", + "top_k":2, "return_documents":false, "truncation":false + }) + ); +} + +#[tokio::test] +async fn duplicate_result_positions_are_not_returned_as_success() { + let (model, _) = scripted(vec![response( + 200, + r#"{ + "data":[{"index":0,"relevance_score":0.9},{"index":0,"relevance_score":0.2}] + }"#, + )]); + let result = model + .rerank(RerankRequest::new("query", vec!["a".into(), "b".into()])) + .await; + assert!(matches!( + result, + Err(Error::Rerank(RerankError::InvalidResponse { .. })) + )); +} + +#[tokio::test] +async fn invalid_results_fail_without_returning_partial_rankings() { + for body in [ + "not json", + "{}", + r#"{"data":[]}"#, + r#"{"data":[{"index":1,"relevance_score":0.5}]}"#, + r#"{"data":[{"index":-1,"relevance_score":0.5}]}"#, + r#"{"data":[{"index":0.5,"relevance_score":0.5}]}"#, + r#"{"data":[{"index":"0","relevance_score":0.5}]}"#, + r#"{"data":[{"relevance_score":0.5}]}"#, + r#"{"data":[{"index":0}]}"#, + r#"{"data":[{"index":0,"relevance_score":"high"}]}"#, + r#"{"data":[{"index":0,"relevance_score":1e400}]}"#, + r#"{"data":[{"index":0,"relevance_score":0.5},{"index":0,"relevance_score":0.3}]}"#, + ] { + let (model, transport) = scripted(vec![response(200, body)]); + assert!( + matches!( + model.rerank(request()).await, + Err(Error::Rerank(RerankError::InvalidResponse { .. })) + ), + "{body}" + ); + assert_eq!(transport.requests.lock().unwrap().len(), 1); + } +} + +#[tokio::test] +async fn results_sort_by_score_and_original_index_without_deduplicating_text() { + let (model, _) = scripted(vec![response( + 200, + r#"{"data":[ + {"index":2,"relevance_score":-2.0}, + {"index":1,"relevance_score":7.0}, + {"index":0,"relevance_score":7.0} + ]}"#, + )]); + let result = model + .rerank(RerankRequest::new("query", vec!["same".into(); 3])) + .await + .unwrap(); + assert_eq!( + result + .results + .iter() + .map(|hit| hit.index) + .collect::>(), + vec![0, 1, 2] + ); + assert_eq!(result.results[2].relevance_score, -2.0); +} + +#[tokio::test] +async fn malformed_usage_does_not_destroy_valid_rankings() { + for usage in [ + json!(null), + json!("bad"), + json!({}), + json!({"total_tokens":-1}), + json!({"total_tokens":1.5}), + json!({"total_tokens":1e30}), + json!({"total_tokens":"7"}), + ] { + let body = json!({"data":[{"index":0,"relevance_score":0.75}], "usage":usage}); + let (model, _) = scripted(vec![response(200, &body.to_string())]); + let result = model.rerank(request()).await.unwrap(); + assert_eq!(result.results.len(), 1); + assert_eq!(result.usage, None); + } + let (model, _) = scripted(vec![response( + 200, + r#"{ + "data":[{"index":0,"relevance_score":0.75}],"usage":{"total_tokens":0} + }"#, + )]); + assert_eq!( + model.rerank(request()).await.unwrap().usage, + Some(RerankUsage { input_tokens: 0 }) + ); +} + +#[tokio::test] +async fn invalid_inputs_fail_before_transport_and_zero_results_are_a_noop() { + for invalid in [ + RerankRequest::new(" \n", vec!["document".into()]), + RerankRequest::new("query", vec!["document".into(), " \t".into()]), + RerankRequest::new("query", vec!["document".into(); 1001]), + ] { + let (model, transport) = scripted(vec![]); + assert!(matches!( + model.rerank(invalid).await, + Err(Error::Validation(_)) + )); + assert!(transport.requests.lock().unwrap().is_empty()); + } + let (model, transport) = scripted(vec![]); + let mut empty = RerankRequest::new("", vec!["".into()]); + empty.top_k = Some(0); + assert!(model.rerank(empty).await.unwrap().results.is_empty()); + assert!(transport.requests.lock().unwrap().is_empty()); +} + +#[test] +fn rejects_invalid_configuration_without_exposing_credentials() { + for endpoint in [ + "", + "file:///tmp/provider", + "https://user:secret@example.com/v1", + "https://example.com/v1?key=secret", + "https://example.com/v1#secret", + "http://example.com/v1", + ] { + let error = VoyageReranker::with_options( + "secret", + VoyageRerankConfig { + base_url: endpoint.into(), + ..VoyageRerankConfig::default() + }, + ) + .unwrap_err(); + assert!(matches!(error, Error::Validation(_))); + assert!(!format!("{error:?}").contains("secret")); + } + for config in [ + VoyageRerankConfig { + model: " ".into(), + ..Default::default() + }, + VoyageRerankConfig { + timeout: Duration::ZERO, + ..Default::default() + }, + VoyageRerankConfig { + max_retries: 4, + ..Default::default() + }, + VoyageRerankConfig { + max_request_bytes: 0, + ..Default::default() + }, + VoyageRerankConfig { + max_response_bytes: 0, + ..Default::default() + }, + ] { + assert!(matches!( + VoyageReranker::with_options("secret", config), + Err(Error::Validation(_)) + )); + } + for key in ["", " \t", "secret\nheader"] { + assert!(matches!( + VoyageReranker::new(key), + Err(Error::Validation(_)) + )); + } +} + +#[tokio::test] +async fn request_limit_counts_json_escaping_and_truncation_is_explicit() { + let (mut model, transport) = scripted(vec![]); + model.config.max_request_bytes = 200; + let too_large = RerankRequest::new("query", vec!["\u{0001}".repeat(40)]); + assert!(matches!( + model.rerank(too_large).await, + Err(Error::Validation(_)) + )); + assert!(transport.requests.lock().unwrap().is_empty()); + + let (model, transport) = scripted(vec![response(200, ONE_RESULT)]); + let mut truncate = request(); + truncate.truncate = true; + model.rerank(truncate).await.unwrap(); + let requests = transport.requests.lock().unwrap(); + let body: Value = + serde_json::from_slice(requests[0].body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(body["truncation"], true); +} + +#[tokio::test] +async fn response_limit_is_enforced_on_chunks_with_absent_or_false_length() { + for length in [None, Some("1"), Some("100")] { + let stream = futures::stream::iter(vec![Ok::<_, std::io::Error>("123"), Ok("456")]); + let mut builder = http::Response::builder(); + if let Some(length) = length { + builder = builder.header("content-length", length); + } + let incoming = builder + .body(reqwest::Body::wrap_stream(stream)) + .unwrap() + .into(); + let (mut model, transport) = scripted(vec![Ok(incoming)]); + model.config.max_response_bytes = 5; + assert!(matches!( + model.rerank(request()).await, + Err(Error::Rerank(RerankError::ResponseTooLarge { limit: 5 })) + )); + assert_eq!(transport.requests.lock().unwrap().len(), 1); + } +} + +#[tokio::test(start_paused = true)] +async fn rate_limit_response_retries_then_returns_success() { + let throttled = http::Response::builder() + .status(429) + .header("retry-after", "2") + .body("secret provider body".to_owned()) + .unwrap() + .into(); + let (model, transport) = scripted(vec![Ok(throttled), response(200, ONE_RESULT)]); + let start = tokio::time::Instant::now(); + assert_eq!(model.rerank(request()).await.unwrap().results[0].index, 0); + assert_eq!(transport.requests.lock().unwrap().len(), 2); + assert_eq!(start.elapsed(), Duration::from_secs(2)); +} + +#[tokio::test] +async fn precancellation_wins_even_for_empty_or_invalid_input() { + for mut input in [ + RerankRequest::new("", vec![]), + RerankRequest::new("", vec!["".into()]), + ] { + let (model, transport) = scripted(vec![]); + input.cancellation.cancel(); + input.top_k = Some(0); + assert!(matches!(model.rerank(input).await, Err(Error::Cancelled))); + assert!(transport.requests.lock().unwrap().is_empty()); + } +} + +#[tokio::test(start_paused = true)] +async fn deadline_covers_a_response_body_that_never_finishes() { + let pending = futures::stream::pending::>(); + let incoming = http::Response::new(reqwest::Body::wrap_stream(pending)).into(); + let (mut model, transport) = scripted(vec![Ok(incoming)]); + model.config.timeout = Duration::from_secs(3); + let result = tokio::time::timeout(Duration::from_secs(10), model.rerank(request())) + .await + .expect("the operation must enforce its own deadline"); + assert!(matches!(result, Err(Error::Rerank(RerankError::Timeout)))); + assert_eq!(transport.requests.lock().unwrap().len(), 1); +} + +#[tokio::test(start_paused = true)] +async fn only_documented_transient_statuses_retry_with_a_bounded_attempt_count() { + for status in [429, 500, 502, 503, 504] { + let (model, transport) = scripted(vec![response(status, ""), response(200, ONE_RESULT)]); + assert!(model.rerank(request()).await.is_ok(), "status {status}"); + assert_eq!(transport.requests.lock().unwrap().len(), 2); + } + let (model, transport) = scripted((0..4).map(|_| response(503, "")).collect()); + assert!(matches!( + model.rerank(request()).await, + Err(Error::Rerank(RerankError::HttpStatus { status: 503 })) + )); + assert_eq!(transport.requests.lock().unwrap().len(), 4); +} + +#[tokio::test] +async fn permanent_failures_and_disabled_retries_make_one_attempt() { + for status in [301, 400, 401, 403, 404, 422, 501] { + let (model, transport) = scripted(vec![response(status, "secret-query test-key")]); + let error = model.rerank(request()).await.unwrap_err(); + match status { + 401 | 403 => assert!( + matches!(error, Error::Rerank(RerankError::Authentication { status: actual }) if actual == status) + ), + _ => assert!( + matches!(error, Error::Rerank(RerankError::HttpStatus { status: actual }) if actual == status) + ), + } + assert!(!format!("{error:?} {error}").contains("secret-query")); + assert!(!format!("{error:?} {error}").contains("test-key")); + assert_eq!(transport.requests.lock().unwrap().len(), 1); + } + let throttled = http::Response::builder() + .status(429) + .header("retry-after", "9") + .body(String::new()) + .unwrap() + .into(); + let (mut model, transport) = scripted(vec![Ok(throttled)]); + model.config.max_retries = 0; + assert!(matches!( + model.rerank(request()).await, + Err(Error::Rerank(RerankError::RateLimited { + retry_after_ms: Some(9000) + })) + )); + assert_eq!(transport.requests.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn transport_and_body_read_failures_are_not_replayed() { + let failing_body = futures::stream::iter(vec![Err::(std::io::Error::other( + "test-key secret-query", + ))]); + for incoming in [ + Err(RerankError::Transport.into()), + Ok(http::Response::new(reqwest::Body::wrap_stream(failing_body)).into()), + ] { + let (model, transport) = scripted(vec![incoming]); + let error = model.rerank(request()).await.unwrap_err(); + assert!(matches!(error, Error::Rerank(RerankError::Transport))); + assert_eq!(transport.requests.lock().unwrap().len(), 1); + assert!(!format!("{error:?} {error}").contains("test-key")); + assert!(!format!("{error:?} {error}").contains("secret-query")); + } +} + +struct DropSignal(Arc); + +impl Drop for DropSignal { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } +} + +struct StalledTransport { + dropped: Arc, +} + +#[async_trait] +impl HttpTransport for StalledTransport { + async fn send(&self, _request: reqwest::Request) -> Result { + let _signal = DropSignal(self.dropped.clone()); + std::future::pending().await + } +} + +#[tokio::test] +async fn cancellation_drops_in_flight_transport_and_body_read() { + for waiting_for_body in [false, true] { + let dropped = Arc::new(AtomicBool::new(false)); + let (mut model, _) = scripted(vec![]); + if waiting_for_body { + let signal = DropSignal(dropped.clone()); + let body = futures::stream::once(async move { + let _signal = signal; + std::future::pending::>().await + }); + model = scripted(vec![Ok(http::Response::new(reqwest::Body::wrap_stream( + body, + )) + .into())]) + .0; + } else { + model.transport = Arc::new(StalledTransport { + dropped: dropped.clone(), + }); + } + let input = request(); + let cancel = input.cancellation.clone(); + let operation = model.rerank(input); + tokio::pin!(operation); + assert!(futures::poll!(&mut operation).is_pending()); + assert!(!dropped.load(Ordering::SeqCst)); + cancel.cancel(); + assert!(matches!(operation.await, Err(Error::Cancelled))); + assert!(dropped.load(Ordering::SeqCst)); + } +} + +#[tokio::test(start_paused = true)] +async fn cancellation_interrupts_pacing_and_backoff_without_later_attempts() { + for waiting_for_pacing in [false, true] { + let (mut model, transport) = scripted(vec![response(429, "")]); + if waiting_for_pacing { + model.config.base_url = "https://rerank-cancel-pacing.invalid/v1".into(); + crate::rate_limit::acquire(&model.config.base_url).await; + } + let input = request(); + let cancel = input.cancellation.clone(); + let operation = model.rerank(input); + tokio::pin!(operation); + assert!(futures::poll!(&mut operation).is_pending()); + cancel.cancel(); + assert!(matches!(operation.await, Err(Error::Cancelled))); + tokio::time::advance(Duration::from_secs(60)).await; + assert_eq!( + transport.requests.lock().unwrap().len(), + usize::from(!waiting_for_pacing) + ); + } +} + +#[tokio::test(start_paused = true)] +async fn total_deadline_covers_pacing_and_all_retries() { + let (mut model, transport) = scripted(vec![]); + model.config.base_url = "https://rerank-deadline-pacing.invalid/v1".into(); + model.config.timeout = Duration::from_millis(100); + crate::rate_limit::acquire(&model.config.base_url).await; + assert!(matches!( + model.rerank(request()).await, + Err(Error::Rerank(RerankError::Timeout)) + )); + assert!(transport.requests.lock().unwrap().is_empty()); + + let (mut model, transport) = scripted(vec![response(503, ""), response(503, "")]); + model.config.timeout = Duration::from_secs(2); + let start = tokio::time::Instant::now(); + assert!(matches!( + model.rerank(request()).await, + Err(Error::Rerank(RerankError::Timeout)) + )); + assert_eq!(transport.requests.lock().unwrap().len(), 2); + assert_eq!(start.elapsed(), Duration::from_secs(2)); +} + +struct CancellingTransport { + cancellation: RerankCancellation, +} + +#[async_trait] +impl HttpTransport for CancellingTransport { + async fn send(&self, _request: reqwest::Request) -> Result { + self.cancellation.cancel(); + Err(RerankError::Transport.into()) + } +} + +#[tokio::test(start_paused = true)] +async fn observed_cancellation_wins_over_failure_and_simultaneous_deadline() { + let (mut model, _) = scripted(vec![]); + let input = request(); + model.transport = Arc::new(CancellingTransport { + cancellation: input.cancellation.clone(), + }); + assert!(matches!(model.rerank(input).await, Err(Error::Cancelled))); + + model.transport = Arc::new(StalledTransport { + dropped: Arc::new(AtomicBool::new(false)), + }); + model.config.timeout = Duration::from_secs(1); + let input = request(); + let cancel = input.cancellation.clone(); + let operation = model.rerank(input); + tokio::pin!(operation); + assert!(futures::poll!(&mut operation).is_pending()); + tokio::time::advance(Duration::from_secs(1)).await; + cancel.cancel(); + assert!(matches!(operation.await, Err(Error::Cancelled))); +} + +#[test] +fn debug_omits_document_content_and_credentials() { + let model = VoyageReranker::with_options( + "secret-key", + VoyageRerankConfig { + model: "secret-key".into(), + base_url: "https://example.com/secret-key".into(), + ..Default::default() + }, + ) + .unwrap(); + assert!(!format!("{model:?}").contains("secret-key")); + let request = RerankRequest::new("private-query", vec!["private-document".into()]); + let debug = format!("{request:?}"); + assert!(!debug.contains("private-query")); + assert!(!debug.contains("private-document")); +} diff --git a/crates/tinyinference-embeddings/src/types.rs b/crates/tinyinference-embeddings/src/types.rs index 5bee5080..789d8e16 100644 --- a/crates/tinyinference-embeddings/src/types.rs +++ b/crates/tinyinference-embeddings/src/types.rs @@ -61,7 +61,7 @@ impl EmbeddingCancellation { self.cancelled.load(Ordering::Acquire) } - async fn cancelled(&self) { + pub(crate) async fn cancelled(&self) { if self.is_cancelled() { return; } diff --git a/crates/tinyinference-voice/src/postprocess.rs b/crates/tinyinference-voice/src/postprocess.rs index 3ab28ebf..efc079bb 100644 --- a/crates/tinyinference-voice/src/postprocess.rs +++ b/crates/tinyinference-voice/src/postprocess.rs @@ -1,6 +1,6 @@ //! Local-LLM transcription cleanup. -use std::time::Duration; +use std::{future::Future, time::Duration}; use tinyinference_local::service::{LocalAiService, RuntimeConfig}; @@ -36,6 +36,14 @@ pub async fn cleanup_transcription( ); let inference = service.inference_interactive(runtime, CLEANUP_SYSTEM_PROMPT, &prompt, Some(512), true); + cleanup_or_original(inference, raw_text, timeout).await +} + +async fn cleanup_or_original( + inference: impl Future>, + raw_text: &str, + timeout: Duration, +) -> String { match tokio::time::timeout(timeout, inference).await { Ok(Ok(cleaned)) if !cleaned.trim().is_empty() => cleaned.trim().to_string(), _ => raw_text.to_string(), diff --git a/crates/tinyinference-voice/src/postprocess_tests.rs b/crates/tinyinference-voice/src/postprocess_tests.rs index 67e579ca..5219f110 100644 --- a/crates/tinyinference-voice/src/postprocess_tests.rs +++ b/crates/tinyinference-voice/src/postprocess_tests.rs @@ -1,8 +1,8 @@ +use super::*; + use axum::{Json, Router, routing::post}; use serde_json::{Value, json}; -use super::*; - const TIMEOUT: Duration = Duration::from_secs(5); async fn spawn_mock(app: Router) -> String { @@ -127,35 +127,49 @@ async fn llm_error_response_falls_back_to_raw_text() { #[tokio::test(start_paused = true)] async fn slow_llm_times_out_and_falls_back_to_raw_text() { - let started = std::sync::Arc::new(tokio::sync::Notify::new()); - let handler_started = started.clone(); - let app = Router::new().route( - "/v1/chat/completions", - post(move || { - let handler_started = handler_started.clone(); - async move { - handler_started.notify_one(); - std::future::pending::>().await - } - }), - ); - let base = spawn_mock(app).await; - let (runtime, service) = runtime_and_service(&base, "ready"); - let cleanup = tokio::spawn(async move { - cleanup_transcription( - &service, - &runtime, - "raw text", - None, - Duration::from_millis(100), - ) - .await - }); - started.notified().await; - tokio::time::advance(Duration::from_millis(101)).await; - let result = cleanup.await.unwrap(); + use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }; + use std::task::Poll; + + struct DropSignal(Arc); + impl Drop for DropSignal { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + + let dropped = Arc::new(AtomicBool::new(false)); + let signal = DropSignal(dropped.clone()); + let inference = async move { + let _signal = signal; + std::future::pending::>().await + }; + let cleanup = cleanup_or_original(inference, "raw text", Duration::from_millis(100)); + tokio::pin!(cleanup); + + std::future::poll_fn(|cx| { + assert!(cleanup.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + tokio::time::advance(Duration::from_millis(99)).await; + std::future::poll_fn(|cx| { + assert!(cleanup.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + assert!(!dropped.load(Ordering::SeqCst)); + tokio::time::advance(Duration::from_millis(2)).await; + let result = std::future::poll_fn(|cx| match cleanup.as_mut().poll(cx) { + Poll::Ready(result) => Poll::Ready(result), + Poll::Pending => panic!("cleanup did not finish at its deadline"), + }) + .await; assert_eq!(result, "raw text"); + assert!(dropped.load(Ordering::SeqCst)); } /// Captures the user prompt the model receives, echoing it back.