pub mod local;
pub mod stub;
use anyhow::Result;
pub use local::LocalKbEmbedder;
pub use stub::StubEmbedder;
pub trait KbEmbedder: Send + Sync {
fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>>;
fn dimension(&self) -> usize;
fn embedder_id(&self) -> &str;
}
pub fn resolve_embedder(kb_root: &std::path::Path) -> std::sync::Arc<dyn KbEmbedder> {
use std::sync::Arc;
let embed_cfg = effective_embed_config();
if let Some(cfg) = embed_cfg.as_ref() {
match cfg.provider.as_deref() {
Some("openai") => {
let model = cfg
.model
.clone()
.unwrap_or_else(|| crate::embed::OPENAI_DEFAULT_MODEL.to_owned());
let api_key = cfg
.api_key
.as_ref()
.and_then(|s| s.resolve_early())
.or_else(|| std::env::var("OPENAI_API_KEY").ok());
let dim = cfg
.dimensions
.unwrap_or_else(|| crate::embed::openai_model_dim(&model))
as usize;
let base_url = cfg
.base_url
.clone()
.unwrap_or_else(|| crate::embed::OPENAI_DEFAULT_BASE_URL.to_owned());
tracing::info!(model = %model, dim, base_url = %base_url, "kb: using remote OpenAI-compatible embedder");
return Arc::new(LocalKbEmbedder::remote_openai(
base_url, model, api_key, dim,
));
}
None | Some("local") => {}
Some(other) => {
tracing::warn!(
provider = other,
"kb: no embedder adapter for this provider; falling back to local BGE scan"
);
}
}
}
let base_dir = kb_root.parent().unwrap_or(kb_root);
let models_dir = base_dir.join("models");
let preferred = embed_cfg
.as_ref()
.and_then(|c| c.local.as_ref())
.and_then(|l| l.model_repo.as_deref())
.and_then(|repo| repo.rsplit('/').next())
.map(str::to_owned);
let defaults = ["bge-small-zh", "bge-base-zh", "bge-small-en"];
let candidates = preferred.iter().map(String::as_str).chain(defaults);
for name in candidates {
let dir = models_dir.join(name);
if dir.join("model.safetensors").exists() {
match LocalKbEmbedder::load(&dir) {
Ok(e) => {
let dim = KbEmbedder::dimension(&e);
tracing::info!(model = name, dim, "kb: using local BGE embedder");
return Arc::new(e);
}
Err(e) => {
tracing::warn!(model = name, "kb: BGE load failed, trying next: {e:#}");
}
}
}
}
tracing::info!("kb: no local BGE model found, using StubEmbedder (1024-dim)");
Arc::new(StubEmbedder::default())
}
pub fn effective_embed_config() -> Option<crate::config::schema::EmbedConfig> {
let cfg = crate::config::load().ok()?;
cfg.raw
.kb
.as_ref()
.and_then(|k| k.embed.clone())
.or_else(|| cfg.raw.memory_search.clone())
}