use std::any::TypeId;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use crate::kit::{AsyncAutoBuilder, AsyncKit, ModuleMeta};
use super::client::{EmbedClient, LocalEmbedClient, OpenAIEmbedClient};
use super::{EmbedError, EmbeddingConfig, Result};
pub type EmbedConfig = EmbeddingConfig;
pub struct EmbedModule;
impl ModuleMeta for EmbedModule {
const NAME: &'static str = "embed";
fn dependencies() -> &'static [(&'static str, TypeId)] {
&[]
}
}
impl AsyncAutoBuilder for EmbedModule {
type Capability = Arc<dyn EmbedClient>;
type Error = EmbedError;
fn build<'a>(
kit: &'a AsyncKit,
) -> Pin<Box<dyn Future<Output = std::result::Result<Self::Capability, Self::Error>> + Send + 'a>>
{
Box::pin(async move {
let config = kit
.config::<EmbeddingConfig>()
.map_err(|e| EmbedError::Unavailable(e.to_string()))?;
Self::build_cap(&config)
})
}
}
impl EmbedModule {
pub(crate) fn build_cap(config: &EmbeddingConfig) -> Result<Arc<dyn EmbedClient>> {
Ok(Arc::new(EmbedCapability {
config: config.clone(),
local_client: Mutex::new(None),
}))
}
}
struct EmbedCapability {
config: EmbeddingConfig,
local_client: Mutex<Option<LocalEmbedClient>>,
}
impl EmbedClient for EmbedCapability {
fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
if self.config.is_local() {
let mut guard = self.local_client.lock().map_err(|e| {
EmbedError::Unavailable(format!("local_client mutex poisoned: {e}"))
})?;
if guard.is_none() {
let client = LocalEmbedClient::new(&self.config)?;
*guard = Some(client);
}
guard
.as_ref()
.expect("local_client initialized")
.embed(texts)
} else {
if !self.config.has_api_key() {
return Err(EmbedError::MissingApiKey);
}
let client = OpenAIEmbedClient::new(self.config.clone())?;
client.embed(texts)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::kit::{AsyncKit, EmbedModule};
#[test]
fn build_returns_send_sync_capability() {
let cap =
EmbedModule::build_cap(&EmbeddingConfig::default()).expect("EmbedModule::build_cap");
fn _assert_send_sync<T: Send + Sync>(_: &T) {}
_assert_send_sync(&cap);
}
#[test]
fn capability_embed_local_without_model_returns_unavailable() {
std::env::remove_var(crate::embed::EMBED_ENDPOINT_ENV);
std::env::remove_var(crate::embed::EMBED_MODEL_PATH_ENV);
let cap =
EmbedModule::build_cap(&EmbeddingConfig::default()).expect("EmbedModule::build_cap");
let result = cap.embed(&["hello"]);
assert!(result.is_err(), "should error without model file");
let err = result.unwrap_err();
assert!(
matches!(err, EmbedError::Unavailable(ref msg) if msg.contains("not found")),
"expected Unavailable with 'not found', got: {err}"
);
}
#[test]
fn capability_embed_remote_without_api_key_returns_missing_api_key() {
std::env::remove_var(crate::embed::API_KEY_ENV);
std::env::remove_var(crate::embed::OPENAI_API_KEY_ENV);
let cap = EmbedModule::build_cap(&EmbeddingConfig {
endpoint: Some("https://api.openai.com/v1".to_string()),
..EmbeddingConfig::default()
})
.expect("EmbedModule::build_cap");
let result = cap.embed(&["hello"]);
assert!(
matches!(result, Err(EmbedError::MissingApiKey)),
"expected MissingApiKey, got {result:?}"
);
}
#[tokio::test]
async fn kit_registration_flow() {
let mut kit = AsyncKit::new();
kit.set_config(EmbeddingConfig::default());
kit.register::<EmbedModule>()
.expect("register::<EmbedModule>");
let kit = kit.build().await.expect("build");
assert!(kit.contains::<EmbedModule>(), "EmbedModule missing");
let _required = kit
.require::<EmbedModule>()
.expect("require::<EmbedModule>");
}
#[test]
fn embed_config_alias_matches_embedding_config() {
let cfg: EmbedConfig = EmbeddingConfig::default();
assert!(cfg.is_local());
}
}