use async_trait::async_trait;
use super::voyageai::{VoyageAIConfig, VoyageAITextVectorizer};
use super::{AsyncVectorizer, Vectorizer};
use crate::error::Result;
const DEFAULT_MODEL: &str = "voyage-3-large";
#[derive(Debug, Clone)]
pub struct AnthropicConfig {
pub(crate) inner: VoyageAIConfig,
}
impl AnthropicConfig {
pub fn new(
api_key: impl Into<String>,
model: impl Into<String>,
input_type: Option<String>,
) -> Self {
Self {
inner: VoyageAIConfig::new(api_key, model, input_type),
}
}
pub fn with_defaults(api_key: impl Into<String>, input_type: Option<String>) -> Self {
Self::new(api_key, DEFAULT_MODEL, input_type)
}
pub fn from_env(input_type: Option<String>) -> Result<Self> {
let inner = VoyageAIConfig::from_env(DEFAULT_MODEL, input_type)?;
Ok(Self { inner })
}
pub fn from_env_with_model(
model: impl Into<String>,
input_type: Option<String>,
) -> Result<Self> {
let inner = VoyageAIConfig::from_env(model, input_type)?;
Ok(Self { inner })
}
}
#[derive(Debug, Clone)]
pub struct AnthropicTextVectorizer {
inner: VoyageAITextVectorizer,
}
impl AnthropicTextVectorizer {
pub fn new(config: AnthropicConfig) -> Self {
Self {
inner: VoyageAITextVectorizer::new(config.inner),
}
}
}
impl Vectorizer for AnthropicTextVectorizer {
fn embed(&self, text: &str) -> Result<Vec<f32>> {
Vectorizer::embed(&self.inner, text)
}
fn embed_many(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
Vectorizer::embed_many(&self.inner, texts)
}
}
#[async_trait]
impl AsyncVectorizer for AnthropicTextVectorizer {
async fn embed(&self, text: &str) -> Result<Vec<f32>> {
AsyncVectorizer::embed(&self.inner, text).await
}
async fn embed_many(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
AsyncVectorizer::embed_many(&self.inner, texts).await
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn anthropic_config_stores_fields() {
let cfg = AnthropicConfig::new("key", "voyage-3-large", Some("document".into()));
assert_eq!(cfg.inner.api_key, "key");
assert_eq!(cfg.inner.model, "voyage-3-large");
assert_eq!(cfg.inner.input_type.as_deref(), Some("document"));
}
#[test]
fn anthropic_config_with_defaults_uses_default_model() {
let cfg = AnthropicConfig::with_defaults("key", None);
assert_eq!(cfg.inner.model, DEFAULT_MODEL);
assert!(cfg.inner.input_type.is_none());
}
#[test]
fn anthropic_vectorizer_is_send_sync() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<AnthropicTextVectorizer>();
}
}