use std::path::PathBuf;
pub const API_KEY_ENV: &str = "CODENEXUS_EMBED_API_KEY";
pub const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
pub const EMBED_ENDPOINT_ENV: &str = "CODENEXUS_EMBED_ENDPOINT";
pub const EMBED_MODEL_PATH_ENV: &str = "CODENEXUS_EMBED_MODEL_PATH";
pub const DEFAULT_MODEL_PATH: &str = "assets/arctic-embed-xs.onnx";
pub const DEFAULT_TOKENIZER_FILENAME: &str = "tokenizer.json";
#[derive(Debug, Clone)]
pub struct EmbeddingConfig {
pub endpoint: Option<String>,
pub model: String,
pub api_key: Option<String>,
pub model_path: Option<PathBuf>,
}
impl Default for EmbeddingConfig {
fn default() -> Self {
Self {
endpoint: None,
model: "arctic-embed-xs".to_string(),
api_key: None,
model_path: None,
}
}
}
impl EmbeddingConfig {
#[must_use]
pub fn from_env() -> Self {
let endpoint = std::env::var(EMBED_ENDPOINT_ENV).ok();
let api_key = std::env::var(API_KEY_ENV)
.or_else(|_| std::env::var(OPENAI_API_KEY_ENV))
.ok();
let model_path = std::env::var(EMBED_MODEL_PATH_ENV).ok().map(PathBuf::from);
Self {
endpoint,
api_key,
model_path,
..Self::default()
}
}
#[must_use]
pub fn is_local(&self) -> bool {
self.endpoint.is_none()
}
#[must_use]
pub fn is_remote(&self) -> bool {
self.endpoint.is_some()
}
#[must_use]
pub fn has_api_key(&self) -> bool {
self.api_key.is_some()
}
#[must_use]
pub fn resolved_model_path(&self) -> PathBuf {
self.model_path
.clone()
.unwrap_or_else(|| PathBuf::from(DEFAULT_MODEL_PATH))
}
#[must_use]
pub fn resolved_tokenizer_path(&self) -> PathBuf {
self.resolved_model_path()
.with_file_name(DEFAULT_TOKENIZER_FILENAME)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn config_default_is_local_offline() {
let cfg = EmbeddingConfig::default();
assert!(cfg.is_local(), "default should be local (offline)");
assert!(!cfg.is_remote(), "default should not be remote");
assert!(cfg.endpoint.is_none(), "default endpoint should be None");
assert!(!cfg.model.is_empty());
assert!(!cfg.has_api_key());
}
#[test]
fn config_default_uses_arctic_model_name() {
let cfg = EmbeddingConfig::default();
assert_eq!(cfg.model, "arctic-embed-xs");
}
#[test]
fn config_default_resolves_default_model_path() {
let cfg = EmbeddingConfig::default();
assert_eq!(cfg.resolved_model_path(), PathBuf::from(DEFAULT_MODEL_PATH));
assert_eq!(
cfg.resolved_tokenizer_path(),
PathBuf::from("assets/tokenizer.json")
);
}
#[test]
fn config_with_custom_model_path() {
let cfg = EmbeddingConfig {
model_path: Some(PathBuf::from("/custom/model.onnx")),
..EmbeddingConfig::default()
};
assert_eq!(
cfg.resolved_model_path(),
PathBuf::from("/custom/model.onnx")
);
assert_eq!(
cfg.resolved_tokenizer_path(),
PathBuf::from("/custom/tokenizer.json")
);
}
#[test]
fn config_remote_mode_when_endpoint_set() {
let cfg = EmbeddingConfig {
endpoint: Some("https://api.openai.com/v1".to_string()),
..EmbeddingConfig::default()
};
assert!(cfg.is_remote());
assert!(!cfg.is_local());
}
#[test]
fn config_from_env_reads_key() {
let _lock = crate::embed::ENV_TEST_LOCK.lock().unwrap();
std::env::set_var(API_KEY_ENV, "test-key-123");
std::env::remove_var(EMBED_ENDPOINT_ENV);
std::env::remove_var(EMBED_MODEL_PATH_ENV);
let cfg = EmbeddingConfig::from_env();
assert_eq!(cfg.api_key.as_deref(), Some("test-key-123"));
assert!(cfg.has_api_key());
std::env::remove_var(API_KEY_ENV);
}
#[test]
fn config_from_env_falls_back_to_openai_var() {
let _lock = crate::embed::ENV_TEST_LOCK.lock().unwrap();
std::env::remove_var(API_KEY_ENV);
std::env::set_var(OPENAI_API_KEY_ENV, "openai-fallback");
std::env::remove_var(EMBED_ENDPOINT_ENV);
let cfg = EmbeddingConfig::from_env();
assert_eq!(cfg.api_key.as_deref(), Some("openai-fallback"));
std::env::remove_var(OPENAI_API_KEY_ENV);
}
#[test]
fn config_from_env_no_key_returns_none() {
let _lock = crate::embed::ENV_TEST_LOCK.lock().unwrap();
std::env::remove_var(API_KEY_ENV);
std::env::remove_var(OPENAI_API_KEY_ENV);
std::env::remove_var(EMBED_ENDPOINT_ENV);
let cfg = EmbeddingConfig::from_env();
assert!(!cfg.has_api_key());
assert!(cfg.api_key.is_none());
}
#[test]
fn config_from_env_reads_endpoint_for_remote_mode() {
let _lock = crate::embed::ENV_TEST_LOCK.lock().unwrap();
std::env::set_var(EMBED_ENDPOINT_ENV, "https://custom.example.com/v1");
std::env::remove_var(API_KEY_ENV);
std::env::remove_var(OPENAI_API_KEY_ENV);
let cfg = EmbeddingConfig::from_env();
assert!(
cfg.is_remote(),
"setting endpoint should enable remote mode"
);
assert_eq!(
cfg.endpoint.as_deref(),
Some("https://custom.example.com/v1")
);
std::env::remove_var(EMBED_ENDPOINT_ENV);
}
#[test]
fn config_from_env_reads_model_path() {
let _lock = crate::embed::ENV_TEST_LOCK.lock().unwrap();
std::env::set_var(EMBED_MODEL_PATH_ENV, "/custom/from/env/model.onnx");
std::env::remove_var(EMBED_ENDPOINT_ENV);
let cfg = EmbeddingConfig::from_env();
assert_eq!(
cfg.resolved_model_path(),
PathBuf::from("/custom/from/env/model.onnx")
);
std::env::remove_var(EMBED_MODEL_PATH_ENV);
}
#[test]
fn config_from_env_defaults_to_local_when_endpoint_unset() {
let _lock = crate::embed::ENV_TEST_LOCK.lock().unwrap();
std::env::remove_var(EMBED_ENDPOINT_ENV);
let cfg = EmbeddingConfig::from_env();
assert!(
cfg.is_local(),
"should default to local when endpoint unset"
);
}
}