From c1ef0cc7951ff79598edc821a37ca4406e19f9e1 Mon Sep 17 00:00:00 2001 From: William Barber Date: Mon, 27 Apr 2026 15:10:05 -0400 Subject: [PATCH 1/2] support voyage contextual embeddings --- README.md | 16 +++ src/client.rs | 86 ++++++++++++++- src/data/models.json | 29 ++++- src/lib.rs | 5 +- src/models.rs | 39 +++++++ src/providers/mod.rs | 15 ++- src/providers/voyageai.rs | 224 +++++++++++++++++++++++++++++++++++++- 7 files changed, 409 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 78f84f4..15f46e2 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,22 @@ let response = client.embed_with_options( ).await?; ``` +### Contextualized Embeddings + +```rust +use catsu::{Client, InputType}; + +let response = client.contextualized_embed_with_options( + "voyage-context-3", + vec![ + vec!["doc 1 chunk 1".to_string(), "doc 1 chunk 2".to_string()], + vec!["doc 2 chunk 1".to_string()], + ], + Some(InputType::Document), + None, +).await?; +``` + ### Model Catalog ```rust diff --git a/src/client.rs b/src/client.rs index 53935be..fe44eaa 100644 --- a/src/client.rs +++ b/src/client.rs @@ -6,7 +6,9 @@ use std::sync::Arc; use crate::catalog::find_model_by_name; use crate::errors::ClientError; use crate::http::{HttpClient, HttpConfig}; -use crate::models::{EmbedRequest, EmbedResponse, InputType}; +use crate::models::{ + ContextualizedEmbedRequest, ContextualizedEmbedResponse, EmbedRequest, EmbedResponse, InputType, +}; use crate::providers::{ CloudflareProvider, CohereProvider, DeepInfraProvider, EmbeddingProvider, GeminiProvider, JinaProvider, MistralProvider, MixedbreadProvider, NomicProvider, OpenAIProvider, @@ -275,6 +277,84 @@ impl Client { provider_impl.embed(request).await } + /// Generate contextualized embeddings for grouped inputs. + /// + /// Each inner input list is embedded together, so chunks in the same group can + /// contribute context to each other. + pub async fn contextualized_embed( + &self, + model: &str, + inputs: Vec>, + ) -> Result { + self.contextualized_embed_with_options(model, inputs, None, None) + .await + } + + /// Generate contextualized embeddings with additional options. + pub async fn contextualized_embed_with_options( + &self, + model: &str, + inputs: Vec>, + input_type: Option, + dimensions: Option, + ) -> Result { + self.contextualized_embed_full(model, inputs, input_type, dimensions, None) + .await + } + + /// Generate contextualized embeddings with all options including explicit provider. + pub async fn contextualized_embed_full( + &self, + model: &str, + inputs: Vec>, + input_type: Option, + dimensions: Option, + provider: Option<&str>, + ) -> Result { + self.contextualized_embed_with_api_key( + model, inputs, input_type, dimensions, provider, None, + ) + .await + } + + /// Generate contextualized embeddings with all options including API key override. + pub async fn contextualized_embed_with_api_key( + &self, + model: &str, + inputs: Vec>, + input_type: Option, + dimensions: Option, + provider: Option<&str>, + api_key: Option, + ) -> Result { + let (provider_name, model_name) = if let Some(p) = provider { + let model_name = model + .split_once(':') + .map(|(_, m)| m.to_string()) + .unwrap_or_else(|| model.to_string()); + (p.to_string(), model_name) + } else { + self.parse_model_string(model)? + }; + + let provider_impl = + self.providers + .get(&provider_name) + .ok_or_else(|| ClientError::ProviderNotFound { + provider: provider_name.clone(), + })?; + + let request = ContextualizedEmbedRequest { + model: model_name, + inputs, + input_type, + dimensions, + api_key, + }; + + provider_impl.contextualized_embed(request).await + } + /// Parse model string into provider and model name. /// /// Supports formats: @@ -349,6 +429,10 @@ mod tests { assert_eq!(provider, "voyageai"); assert_eq!(model, "voyage-3-large"); + let (provider, model) = client.parse_model_string("voyage-context-3").unwrap(); + assert_eq!(provider, "voyageai"); + assert_eq!(model, "voyage-context-3"); + // Unknown model without prefix - should error let result = client.parse_model_string("non-existent-model-xyz"); assert!(result.is_err()); diff --git a/src/data/models.json b/src/data/models.json index f770d89..254d7b9 100644 --- a/src/data/models.json +++ b/src/data/models.json @@ -1347,6 +1347,33 @@ ], "release_date": "2026-01-15" }, + { + "name": "voyage-context-3", + "provider": "voyageai", + "dimensions": 1024, + "max_input_tokens": 32000, + "cost_per_million_tokens": 0.18, + "modalities": [ + "text" + ], + "supports_batching": true, + "supports_input_type": true, + "supports_dimensions": true, + "mteb_score": null, + "rteb_score": null, + "tokenizer": { + "engine": "huggingface", + "name": "voyageai/voyage-context-3" + }, + "quantizations": [ + "float", + "int8", + "uint8", + "binary", + "ubinary" + ], + "release_date": "2025-07-23" + }, { "name": "voyage-3-large", "provider": "voyageai", @@ -1604,4 +1631,4 @@ "release_date": "2024-11-12" } ] -} \ No newline at end of file +} diff --git a/src/lib.rs b/src/lib.rs index 8463d18..d2bc6cf 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -60,5 +60,8 @@ pub use catalog::{find_model_by_name, get_model, list_catalog_providers, list_mo pub use client::Client; pub use errors::ClientError; pub use http::{HttpClient, HttpConfig}; -pub use models::{EmbedRequest, EmbedResponse, InputType, ModelInfo, Usage}; +pub use models::{ + ContextualizedEmbedRequest, ContextualizedEmbedResponse, EmbedRequest, EmbedResponse, + InputType, ModelInfo, Usage, +}; pub use providers::EmbeddingProvider; diff --git a/src/models.rs b/src/models.rs index 358d02f..a912cba 100644 --- a/src/models.rs +++ b/src/models.rs @@ -23,6 +23,29 @@ pub struct EmbedResponse { pub usage: Usage, } +/// Response from a contextualized embedding request. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ContextualizedEmbedResponse { + /// The generated embeddings grouped by input document/query group. + pub embeddings: Vec>>, + /// The model used to generate embeddings. + pub model: String, + /// The provider that generated the embeddings. + pub provider: String, + /// The dimensionality of the embeddings. + pub dimensions: usize, + /// Number of input groups that were embedded. + pub input_count: usize, + /// Total number of chunks embedded across all input groups. + pub chunk_count: usize, + /// The input type used (if specified). + pub input_type: Option, + /// Latency of the request in milliseconds. + pub latency_ms: f64, + /// Usage information. + pub usage: Usage, +} + /// Token and cost usage information. #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct Usage { @@ -48,6 +71,22 @@ pub struct EmbedRequest { pub api_key: Option, } +/// Request for generating contextualized embeddings. +#[derive(Debug, Clone, Serialize)] +pub struct ContextualizedEmbedRequest { + /// The model to use. + pub model: String, + /// Input text groups to embed together. + pub inputs: Vec>, + /// Optional input type hint (query or document). + pub input_type: Option, + /// Optional output dimensions (if model supports it). + pub dimensions: Option, + /// Optional API key override for this request. + #[serde(skip)] + pub api_key: Option, +} + /// Input type hint for embeddings. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] diff --git a/src/providers/mod.rs b/src/providers/mod.rs index ccf8562..0ef49fb 100644 --- a/src/providers/mod.rs +++ b/src/providers/mod.rs @@ -15,7 +15,9 @@ pub mod voyageai; use async_trait::async_trait; use crate::errors::ClientError; -use crate::models::{EmbedRequest, EmbedResponse}; +use crate::models::{ + ContextualizedEmbedRequest, ContextualizedEmbedResponse, EmbedRequest, EmbedResponse, +}; pub use cloudflare::CloudflareProvider; pub use cohere::CohereProvider; @@ -37,4 +39,15 @@ pub trait EmbeddingProvider: Send + Sync { /// Generate embeddings for the given request. async fn embed(&self, request: EmbedRequest) -> Result; + + /// Generate contextualized embeddings for grouped inputs. + async fn contextualized_embed( + &self, + _request: ContextualizedEmbedRequest, + ) -> Result { + Err(ClientError::InvalidInput(format!( + "{} does not support contextualized embeddings", + self.name() + ))) + } } diff --git a/src/providers/voyageai.rs b/src/providers/voyageai.rs index 879d3af..409061f 100644 --- a/src/providers/voyageai.rs +++ b/src/providers/voyageai.rs @@ -8,10 +8,16 @@ use tracing::debug; use crate::errors::ClientError; use crate::http::HttpClient; -use crate::models::{EmbedRequest, EmbedResponse, InputType, Usage}; +use crate::models::{ + ContextualizedEmbedRequest, ContextualizedEmbedResponse, EmbedRequest, EmbedResponse, + InputType, Usage, +}; use crate::providers::EmbeddingProvider; const VOYAGEAI_API_URL: &str = "https://api.voyageai.com/v1/embeddings"; +const VOYAGEAI_CONTEXTUALIZED_API_URL: &str = + "https://api.voyageai.com/v1/contextualizedembeddings"; +const VOYAGEAI_CONTEXTUAL_MODEL: &str = "voyage-context-3"; /// VoyageAI embeddings provider. #[derive(Debug)] @@ -49,6 +55,17 @@ struct VoyageAIEmbeddingRequest<'a> { output_dimension: Option, } +/// VoyageAI contextualized embedding API request body. +#[derive(Debug, Serialize)] +struct VoyageAIContextualizedEmbeddingRequest<'a> { + model: &'a str, + inputs: &'a [Vec], + #[serde(skip_serializing_if = "Option::is_none")] + input_type: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + output_dimension: Option, +} + /// VoyageAI API response. #[derive(Debug, Deserialize)] struct VoyageAIEmbeddingResponse { @@ -57,6 +74,21 @@ struct VoyageAIEmbeddingResponse { usage: VoyageAIUsage, } +/// VoyageAI contextualized embedding API response. +#[derive(Debug, Deserialize)] +struct VoyageAIContextualizedEmbeddingResponse { + data: Vec, + model: String, + usage: VoyageAIUsage, +} + +#[derive(Debug, Deserialize)] +struct VoyageAIContextualizedEmbeddingGroup { + data: Vec, + #[serde(default)] + index: Option, +} + #[derive(Debug, Deserialize)] struct VoyageAIEmbedding { embedding: Vec, @@ -81,6 +113,33 @@ impl EmbeddingProvider for VoyageAIProvider { } async fn embed(&self, request: EmbedRequest) -> Result { + if request.model == VOYAGEAI_CONTEXTUAL_MODEL { + let input_count = request.inputs.len(); + let context_request = ContextualizedEmbedRequest { + model: request.model, + inputs: request + .inputs + .into_iter() + .map(|input| vec![input]) + .collect(), + input_type: request.input_type, + dimensions: request.dimensions, + api_key: request.api_key, + }; + let response = self.contextualized_embed(context_request).await?; + + return Ok(EmbedResponse { + embeddings: response.embeddings.into_iter().flatten().collect(), + model: response.model, + provider: response.provider, + dimensions: response.dimensions, + input_count, + input_type: response.input_type, + latency_ms: response.latency_ms, + usage: response.usage, + }); + } + debug!( model = %request.model, inputs = request.inputs.len(), @@ -164,10 +223,112 @@ impl EmbeddingProvider for VoyageAIProvider { }, }) } + + async fn contextualized_embed( + &self, + request: ContextualizedEmbedRequest, + ) -> Result { + debug!( + model = %request.model, + input_groups = request.inputs.len(), + chunks = request.inputs.iter().map(Vec::len).sum::(), + "Sending VoyageAI contextualized embedding request" + ); + + let input_type_value = request.input_type; + let input_type_str = input_type_to_str(input_type_value); + + let input_count = request.inputs.len(); + let chunk_count = request.inputs.iter().map(Vec::len).sum(); + let body = VoyageAIContextualizedEmbeddingRequest { + model: &request.model, + inputs: &request.inputs, + input_type: input_type_str, + output_dimension: request.dimensions, + }; + + let body_json = serde_json::to_string(&body)?; + let api_key = request.api_key.as_ref().unwrap_or(&self.api_key).clone(); + + let start = Instant::now(); + let response = self + .http_client + .send_with_retry(|client| { + client + .post(VOYAGEAI_CONTEXTUALIZED_API_URL) + .header("Authorization", format!("Bearer {}", api_key)) + .header("Content-Type", "application/json") + .body(body_json.clone()) + }) + .await?; + + let status = response.status().as_u16(); + let response_text = response.text().await?; + let latency_ms = start.elapsed().as_secs_f64() * 1000.0; + + if status != 200 { + if let Ok(error_response) = + serde_json::from_str::(&response_text) + { + return Err(ClientError::Api { + status, + message: error_response.detail, + }); + } + return Err(ClientError::Api { + status, + message: response_text, + }); + } + + let voyage_response: VoyageAIContextualizedEmbeddingResponse = + serde_json::from_str(&response_text)?; + + let mut groups: Vec<_> = voyage_response.data.into_iter().enumerate().collect(); + groups.sort_by_key(|(position, group)| group.index.unwrap_or(*position)); + + let embedding_groups: Vec>> = groups + .into_iter() + .map(|(_, group)| { + let mut embeddings = group.data; + embeddings.sort_by_key(|e| e.index); + embeddings.into_iter().map(|e| e.embedding).collect() + }) + .collect(); + + let dimensions = embedding_groups + .iter() + .find_map(|group| group.first()) + .map_or(0, Vec::len); + let cost = calculate_cost(&request.model, voyage_response.usage.total_tokens); + + Ok(ContextualizedEmbedResponse { + embeddings: embedding_groups, + model: voyage_response.model, + provider: "voyageai".to_string(), + dimensions, + input_count, + chunk_count, + input_type: input_type_value, + latency_ms, + usage: Usage { + tokens: voyage_response.usage.total_tokens, + cost, + }, + }) + } +} + +fn input_type_to_str(input_type: Option) -> Option<&'static str> { + input_type.map(|it| match it { + InputType::Query => "query", + InputType::Document => "document", + }) } fn calculate_cost(model: &str, tokens: u64) -> Option { let price_per_million = match model { + m if m.contains("voyage-context-3") => 0.18, m if m.contains("voyage-3-lite") => 0.02, m if m.contains("voyage-3") => 0.06, m if m.contains("voyage-code-3") => 0.18, @@ -179,3 +340,64 @@ fn calculate_cost(model: &str, tokens: u64) -> Option { Some((tokens as f64 / 1_000_000.0) * price_per_million) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn serializes_contextualized_embedding_request() { + let inputs = vec![ + vec!["doc 1 chunk 1".to_string(), "doc 1 chunk 2".to_string()], + vec!["doc 2 chunk 1".to_string()], + ]; + let body = VoyageAIContextualizedEmbeddingRequest { + model: VOYAGEAI_CONTEXTUAL_MODEL, + inputs: &inputs, + input_type: Some("document"), + output_dimension: Some(512), + }; + + let json = serde_json::to_value(body).unwrap(); + assert_eq!(json["model"], "voyage-context-3"); + assert_eq!(json["input_type"], "document"); + assert_eq!(json["output_dimension"], 512); + assert_eq!(json["inputs"][0][1], "doc 1 chunk 2"); + } + + #[test] + fn deserializes_contextualized_embedding_response() { + let response: VoyageAIContextualizedEmbeddingResponse = serde_json::from_str( + r#"{ + "data": [ + { + "index": 0, + "data": [ + {"index": 1, "embedding": [0.3, 0.4]}, + {"index": 0, "embedding": [0.1, 0.2]} + ] + }, + { + "index": 1, + "data": [ + {"index": 0, "embedding": [0.5, 0.6]} + ] + } + ], + "model": "voyage-context-3", + "usage": {"total_tokens": 42} + }"#, + ) + .unwrap(); + + assert_eq!(response.model, "voyage-context-3"); + assert_eq!(response.usage.total_tokens, 42); + assert_eq!(response.data.len(), 2); + assert_eq!(response.data[0].data.len(), 2); + } + + #[test] + fn prices_contextualized_embeddings() { + assert_eq!(calculate_cost("voyage-context-3", 1_000_000), Some(0.18)); + } +} From 27049f9538a289b2cdae4540150f99e15362a152 Mon Sep 17 00:00:00 2001 From: William Barber Date: Mon, 27 Apr 2026 15:10:12 -0400 Subject: [PATCH 2/2] expose contextual embeddings in python --- packages/python/README_PYPI.md | 15 +++ packages/python/catsu/__init__.py | 12 +- packages/python/src/lib.rs | 186 ++++++++++++++++++++++----- packages/python/tests/test_client.py | 17 +++ 4 files changed, 198 insertions(+), 32 deletions(-) diff --git a/packages/python/README_PYPI.md b/packages/python/README_PYPI.md index e0fbe68..0e7a92f 100644 --- a/packages/python/README_PYPI.md +++ b/packages/python/README_PYPI.md @@ -83,6 +83,21 @@ response = client.embed( ) ``` +## Contextualized Embeddings + +```python +response = client.contextualized_embed( + "voyage-context-3", + [ + ["doc 1 chunk 1", "doc 1 chunk 2"], + ["doc 2 chunk 1"], + ], + input_type="document", +) + +print(response.embeddings[0][0][:5]) +``` + ## Model Catalog ```python diff --git a/packages/python/catsu/__init__.py b/packages/python/catsu/__init__.py index 74d1f2a..17a6173 100644 --- a/packages/python/catsu/__init__.py +++ b/packages/python/catsu/__init__.py @@ -34,7 +34,13 @@ >>> print(arr.shape) # (1, 1536) """ -from catsu._catsu import Client, EmbedResponse, ModelInfo, Usage - -__all__ = ["Client", "EmbedResponse", "ModelInfo", "Usage"] +from catsu._catsu import Client, ContextualizedEmbedResponse, EmbedResponse, ModelInfo, Usage + +__all__ = [ + "Client", + "ContextualizedEmbedResponse", + "EmbedResponse", + "ModelInfo", + "Usage", +] __version__ = "0.1.8" diff --git a/packages/python/src/lib.rs b/packages/python/src/lib.rs index b63301a..3e28d41 100644 --- a/packages/python/src/lib.rs +++ b/packages/python/src/lib.rs @@ -11,8 +11,8 @@ use pyo3_async_runtimes::tokio::future_into_py; use tokio::runtime::Runtime; use catsu::{ - Client as RustClient, EmbedResponse as RustEmbedResponse, HttpConfig, InputType, - ModelInfo as RustModelInfo, + Client as RustClient, ContextualizedEmbedResponse as RustContextualizedEmbedResponse, + EmbedResponse as RustEmbedResponse, HttpConfig, InputType, ModelInfo as RustModelInfo, }; /// Global tokio runtime for async operations. @@ -91,6 +91,30 @@ pub struct EmbedResponse { pub usage: Usage, } +/// Response from a contextualized embedding request. +#[pyclass] +#[derive(Clone)] +pub struct ContextualizedEmbedResponse { + #[pyo3(get)] + pub embeddings: Vec>>, + #[pyo3(get)] + pub model: String, + #[pyo3(get)] + pub provider: String, + #[pyo3(get)] + pub dimensions: usize, + #[pyo3(get)] + pub input_count: usize, + #[pyo3(get)] + pub chunk_count: usize, + #[pyo3(get)] + pub input_type: Option, + #[pyo3(get)] + pub latency_ms: f64, + #[pyo3(get)] + pub usage: Usage, +} + #[pymethods] impl EmbedResponse { /// Convert embeddings to a numpy array. @@ -111,17 +135,33 @@ impl EmbedResponse { impl From for EmbedResponse { fn from(response: RustEmbedResponse) -> Self { - let input_type = response.input_type.map(|it| match it { - InputType::Query => "query".to_string(), - InputType::Document => "document".to_string(), - }); - EmbedResponse { embeddings: response.embeddings, model: response.model, provider: response.provider, dimensions: response.dimensions, input_count: response.input_count, + input_type: input_type_to_string(response.input_type), + latency_ms: response.latency_ms, + usage: Usage { + tokens: response.usage.tokens, + cost: response.usage.cost, + }, + } + } +} + +impl From for ContextualizedEmbedResponse { + fn from(response: RustContextualizedEmbedResponse) -> Self { + let input_type = input_type_to_string(response.input_type); + + ContextualizedEmbedResponse { + embeddings: response.embeddings, + model: response.model, + provider: response.provider, + dimensions: response.dimensions, + input_count: response.input_count, + chunk_count: response.chunk_count, input_type, latency_ms: response.latency_ms, usage: Usage { @@ -132,6 +172,25 @@ impl From for EmbedResponse { } } +fn parse_input_type(input_type: Option<&str>) -> PyResult> { + match input_type { + Some("query") => Ok(Some(InputType::Query)), + Some("document") => Ok(Some(InputType::Document)), + Some(other) => Err(PyRuntimeError::new_err(format!( + "Invalid input_type '{}'. Must be 'query' or 'document'", + other + ))), + None => Ok(None), + } +} + +fn input_type_to_string(input_type: Option) -> Option { + input_type.map(|it| match it { + InputType::Query => "query".to_string(), + InputType::Document => "document".to_string(), + }) +} + /// High-performance embeddings client. /// /// Example: @@ -234,17 +293,7 @@ impl CatsuClient { let model = model.to_string(); let provider_owned = provider.map(|s| s.to_string()); - let input_type_enum = match input_type { - Some("query") => Some(InputType::Query), - Some("document") => Some(InputType::Document), - Some(other) => { - return Err(PyRuntimeError::new_err(format!( - "Invalid input_type '{}'. Must be 'query' or 'document'", - other - ))) - } - None => None, - }; + let input_type_enum = parse_input_type(input_type)?; let rt = GLOBAL_RUNTIME.clone(); let client = &self.inner; @@ -313,17 +362,7 @@ impl CatsuClient { let model = model.to_string(); let provider_owned = provider.map(|s| s.to_string()); - let input_type_enum = match input_type { - Some("query") => Some(InputType::Query), - Some("document") => Some(InputType::Document), - Some(other) => { - return Err(PyRuntimeError::new_err(format!( - "Invalid input_type '{}'. Must be 'query' or 'document'", - other - ))) - } - None => None, - }; + let input_type_enum = parse_input_type(input_type)?; let client = Arc::clone(&self.inner); @@ -344,6 +383,94 @@ impl CatsuClient { }) } + /// Generate contextualized embeddings for grouped input text. + /// + /// Args: + /// model: Model name (with or without "provider:" prefix) + /// inputs: List of text groups. Each inner list is embedded together. + /// provider: Optional explicit provider name (e.g., "voyageai") + /// input_type: Optional input type hint ("query" or "document") + /// dimensions: Optional output dimensions + /// api_key: Optional API key override for this request + /// + /// Example: + /// >>> response = client.contextualized_embed( + /// ... "voyage-context-3", + /// ... [["doc 1 chunk 1", "doc 1 chunk 2"], ["doc 2 chunk 1"]], + /// ... input_type="document", + /// ... ) + #[pyo3(signature = (model, inputs, provider=None, input_type=None, dimensions=None, api_key=None))] + pub fn contextualized_embed( + &self, + py: Python<'_>, + model: &str, + inputs: Vec>, + provider: Option<&str>, + input_type: Option<&str>, + dimensions: Option, + api_key: Option, + ) -> PyResult { + let model = model.to_string(); + let provider_owned = provider.map(|s| s.to_string()); + let input_type_enum = parse_input_type(input_type)?; + + let rt = GLOBAL_RUNTIME.clone(); + let client = &self.inner; + + let result = py.allow_threads(move || { + rt.block_on(async { + client + .contextualized_embed_with_api_key( + &model, + inputs, + input_type_enum, + dimensions, + provider_owned.as_deref(), + api_key, + ) + .await + }) + }); + + result + .map(ContextualizedEmbedResponse::from) + .map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + /// Async version of contextualized_embed(). + #[pyo3(signature = (model, inputs, provider=None, input_type=None, dimensions=None, api_key=None))] + pub fn acontextualized_embed<'py>( + &self, + py: Python<'py>, + model: &str, + inputs: Vec>, + provider: Option<&str>, + input_type: Option<&str>, + dimensions: Option, + api_key: Option, + ) -> PyResult> { + let model = model.to_string(); + let provider_owned = provider.map(|s| s.to_string()); + let input_type_enum = parse_input_type(input_type)?; + let client = Arc::clone(&self.inner); + + future_into_py(py, async move { + let result = client + .contextualized_embed_with_api_key( + &model, + inputs, + input_type_enum, + dimensions, + provider_owned.as_deref(), + api_key, + ) + .await + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; + + Ok(ContextualizedEmbedResponse::from(result)) + }) + } + /// List available providers. /// /// Returns: @@ -438,6 +565,7 @@ impl CatsuClient { fn _catsu(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; Ok(()) diff --git a/packages/python/tests/test_client.py b/packages/python/tests/test_client.py index 8c130b7..cf0e06a 100644 --- a/packages/python/tests/test_client.py +++ b/packages/python/tests/test_client.py @@ -62,3 +62,20 @@ def test_client_with_proxy_and_other_options(): timeout=60, ) assert client is not None + + +def test_contextualized_response_is_exported(): + """Test that the contextualized response type is exported.""" + from catsu import ContextualizedEmbedResponse + + assert ContextualizedEmbedResponse is not None + + +def test_voyage_context_model_in_catalog(): + """Test that voyage-context-3 can be auto-detected as a VoyageAI model.""" + from catsu import Client + + client = Client() + models = client.list_models("voyageai") + model_names = {model.name for model in models} + assert "voyage-context-3" in model_names