use anyhow::Context as _;
use oxibrain_core::extraction::ExtractMechanism;
use oxibrain_ports::{LlmPort, TokenizerPort};
use std::sync::Arc;
use crate::cmd::foundation::{
self, FoundationError, ProfileRole, ProviderKind, ProviderProfile, ResolvedProfiles,
SecretResolver, default_secret_resolver,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Provider {
Anthropic,
OpenAi,
Local,
}
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub enum ResolutionSource {
ExplicitOverride(Provider),
FoundationProfile {
profile_id: String,
provider: ProviderKind,
model_id: String,
mechanism: ExtractMechanism,
},
CompatEnv {
kind: ProviderKind,
model_id: String,
},
Local,
}
pub struct ProviderLlm {
pub port: Arc<dyn LlmPort>,
pub model_id: String,
pub mechanism: ExtractMechanism,
pub model_digest: Option<String>,
pub tokenizer: Option<Arc<dyn TokenizerPort>>,
pub source: ResolutionSource,
}
impl ProviderLlm {
pub fn profile_id(&self) -> Option<String> {
match &self.source {
ResolutionSource::FoundationProfile { profile_id, .. } => Some(profile_id.clone()),
ResolutionSource::ExplicitOverride(_)
| ResolutionSource::CompatEnv { .. }
| ResolutionSource::Local => None,
}
}
}
pub fn resolve_provider(
explicit: Option<&str>,
anthropic_key_present: bool,
openai_key_present: bool,
) -> anyhow::Result<Provider> {
match explicit {
Some("anthropic") => Ok(Provider::Anthropic),
Some("openai") => Ok(Provider::OpenAi),
Some("local") => Ok(Provider::Local),
Some(other) => anyhow::bail!(
"unknown OXIBRAIN_LLM_PROVIDER={other} (expected: anthropic|openai|local)"
),
None if anthropic_key_present => Ok(Provider::Anthropic),
None if openai_key_present => Ok(Provider::OpenAi),
None => Ok(Provider::Local),
}
}
pub fn resolve_role() -> ProfileRole {
if let Ok(raw) = std::env::var("OXIBRAIN_LLM_ROLE") {
if let Some(role) = ProfileRole::parse(&raw) {
return role;
}
tracing::warn!(
role = %raw,
"OXIBRAIN_LLM_ROLE is not a known role; falling back to memory.extract"
);
}
ProfileRole::MemoryExtract
}
pub async fn from_env() -> anyhow::Result<ProviderLlm> {
from_env_for_role(resolve_role()).await
}
pub async fn from_env_for_role(role: ProfileRole) -> anyhow::Result<ProviderLlm> {
let explicit = std::env::var("OXIBRAIN_LLM_PROVIDER").ok();
let anthropic_key_present = std::env::var("ANTHROPIC_API_KEY").is_ok();
let openai_key_present = std::env::var("OPENAI_API_KEY").is_ok();
if let Some(name) = explicit.as_deref() {
match resolve_provider(Some(name), anthropic_key_present, openai_key_present)? {
Provider::Anthropic => return anthropic_from_env(),
Provider::OpenAi => return openai_from_env(),
Provider::Local => return local_from_manifest().await,
}
}
let resolved_profiles =
foundation::load_profiles(&foundation::foundation_home()).map_err(anyhow::Error::msg)?;
if let Some(profiles) = resolved_profiles {
if let Some(provider) =
try_foundation_profile(&profiles, role, default_secret_resolver().as_ref()).await?
{
return Ok(provider);
}
}
if anthropic_key_present {
return anthropic_from_env();
}
if openai_key_present {
return openai_from_env();
}
local_from_manifest().await
}
#[doc(hidden)]
pub async fn try_foundation_profile(
profiles: &ResolvedProfiles,
role: ProfileRole,
secret_resolver: &dyn SecretResolver,
) -> anyhow::Result<Option<ProviderLlm>> {
let mut selected_profile: Option<&ProviderProfile> = None;
for profile in profiles.iter() {
if !profile.roles.contains(&role) {
continue;
}
let mechanism = match ProviderKind::parse(&profile.provider) {
Some(ProviderKind::OpenAi) => ExtractMechanism::JsonSchema,
Some(ProviderKind::Anthropic) | None => ExtractMechanism::ToolCall,
};
if !profile.capabilities.clone().satisfies(mechanism) {
anyhow::bail!(
"Foundation profile `{}` rejected: declared capabilities do not satisfy extraction mechanism {:?}",
profile.id,
mechanism
);
}
selected_profile = Some(profile);
break;
}
let profile = match selected_profile {
Some(p) => p,
None => return Ok(None),
};
let mechanism = match ProviderKind::parse(&profile.provider) {
Some(ProviderKind::OpenAi) => ExtractMechanism::JsonSchema,
_ => ExtractMechanism::ToolCall,
};
let secret = match secret_resolver.resolve(&profile.credential) {
Ok(s) => s,
Err(e @ FoundationError::SecretUnavailable { .. }) => {
tracing::warn!("{e}");
return Ok(None);
}
Err(other) => return Err(anyhow::Error::msg(other.to_string())),
};
let provider_kind = ProviderKind::parse(&profile.provider).ok_or_else(|| {
anyhow::anyhow!(
"Foundation profile `{}` has unknown provider `{}`",
profile.id,
profile.provider
)
})?;
let port: Arc<dyn LlmPort> = match provider_kind {
ProviderKind::Anthropic => Arc::new(oxibrain_llm_http::AnthropicLlm::new(
secret,
profile.model.clone(),
)),
ProviderKind::OpenAi => Arc::new(oxibrain_llm_http::OpenAiLlm::new(
secret,
profile.model.clone(),
)),
};
Ok(Some(ProviderLlm {
port,
model_id: profile.model.clone(),
mechanism,
model_digest: None,
tokenizer: None,
source: ResolutionSource::FoundationProfile {
profile_id: profile.id.clone(),
provider: provider_kind,
model_id: profile.model.clone(),
mechanism,
},
}))
}
fn anthropic_from_env() -> anyhow::Result<ProviderLlm> {
let key = std::env::var("ANTHROPIC_API_KEY")
.map_err(|_| anyhow::anyhow!("ANTHROPIC_API_KEY not set (required for extraction)"))?;
let model = std::env::var("ANTHROPIC_MODEL")
.or_else(|_| std::env::var("OXIBRAIN_MODEL"))
.unwrap_or_else(|_| "claude-sonnet-4-5".to_string());
Ok(ProviderLlm {
port: Arc::new(oxibrain_llm_http::AnthropicLlm::new(key, model.clone())),
model_id: model.clone(),
mechanism: ExtractMechanism::ToolCall,
model_digest: None,
tokenizer: None,
source: ResolutionSource::CompatEnv {
kind: ProviderKind::Anthropic,
model_id: model,
},
})
}
fn openai_from_env() -> anyhow::Result<ProviderLlm> {
let key = std::env::var("OPENAI_API_KEY")
.map_err(|_| anyhow::anyhow!("OPENAI_API_KEY not set (required for extraction)"))?;
let model = std::env::var("OPENAI_MODEL")
.or_else(|_| std::env::var("OXIBRAIN_MODEL"))
.unwrap_or_else(|_| "gpt-4o".to_string());
Ok(ProviderLlm {
port: Arc::new(oxibrain_llm_http::OpenAiLlm::new(key, model.clone())),
model_id: model.clone(),
mechanism: ExtractMechanism::JsonSchema,
model_digest: None,
tokenizer: None,
source: ResolutionSource::CompatEnv {
kind: ProviderKind::OpenAi,
model_id: model,
},
})
}
fn extract_entry(
entries: &[oxibrain::models::ModelEntry],
) -> Option<&oxibrain::models::ModelEntry> {
entries
.iter()
.find(|e| e.role == oxibrain::models::ModelRole::Extract)
}
async fn ensure_local_model_present() -> anyhow::Result<()> {
use oxibrain::models::{default_manifest, load_manifest, model_dir, pull_entry, save_manifest};
use oxibrain::pull_plan::{ExtractPullPlan, plan_extract_pull};
let dir = model_dir();
std::fs::create_dir_all(&dir)?;
let manifest = load_manifest().map_err(|e| anyhow::anyhow!("load model manifest: {e}"))?;
let defaults = default_manifest();
let plan = plan_extract_pull(&manifest, &dir, &defaults);
let entry = match plan {
ExtractPullPlan::NoOp => return Ok(()),
ExtractPullPlan::NeedsPullFromManifest(e) => e,
ExtractPullPlan::NeedsBootstrap(e) => {
let mut next = manifest.clone();
if !next.iter().any(|m| m.name == e.name) {
next.push(e.clone());
save_manifest(&next)?;
}
e
}
};
println!(
"pulling local extract model {} ({} MiB) — first use only...",
entry.name, entry.size_mb
);
pull_entry(&entry, &dir, oxibrain::models::cli_progress)
.await
.map_err(|e| anyhow::anyhow!("pull {}: {e}", entry.name))?;
println!(" verified");
Ok(())
}
async fn local_from_manifest() -> anyhow::Result<ProviderLlm> {
use oxibrain::models::{load_manifest, model_dir, verify_entry};
ensure_local_model_present().await?;
let manifest = load_manifest().context("load model manifest")?;
let entry = extract_entry(&manifest)
.ok_or_else(|| anyhow::anyhow!("local extract model could not be resolved after pull"))?;
let dir = model_dir();
verify_entry(entry, &dir)
.map_err(|e| anyhow::anyhow!("model digest mismatch for {}: {e}", entry.name))?;
let path = dir.join(&entry.file);
let llm = Arc::new(
oxibrain_llm_local::LocalLlm::open(&path, oxibrain_llm_local::LocalLlmOptions::default())
.map_err(|e| anyhow::anyhow!("open local model {}: {e}", path.display()))?,
);
Ok(ProviderLlm {
model_id: entry.name.clone(),
mechanism: ExtractMechanism::Grammar,
model_digest: Some(entry.digest.clone()),
port: llm.clone(),
tokenizer: Some(llm),
source: ResolutionSource::Local,
})
}
pub fn config(
model_id: String,
mechanism: ExtractMechanism,
model_digest: Option<String>,
provider_profile_id: Option<String>,
) -> oxibrain_core::extraction::ExtractorConfig {
use oxibrain_core::registry::CORE_V1_MAJOR;
oxibrain_core::extraction::ExtractorConfig {
model_id,
prompt_version: 2, registry_major: CORE_V1_MAJOR,
mechanism,
max_tokens: 8192,
model_digest,
provider_profile_id,
}
}
#[cfg(test)]
mod tests {
use super::*;
static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
#[test]
fn explicit_provider_wins() {
assert_eq!(
resolve_provider(Some("local"), true, true).unwrap(),
Provider::Local
);
assert_eq!(
resolve_provider(Some("openai"), true, false).unwrap(),
Provider::OpenAi
);
assert_eq!(
resolve_provider(Some("anthropic"), false, false).unwrap(),
Provider::Anthropic
);
}
#[test]
fn unknown_provider_is_rejected() {
assert!(resolve_provider(Some("gemini"), false, false).is_err());
}
#[test]
fn no_explicit_and_no_key_falls_back_to_local() {
assert_eq!(
resolve_provider(None, false, false).unwrap(),
Provider::Local
);
}
#[test]
fn anthropic_key_preferred_over_local() {
assert_eq!(
resolve_provider(None, true, false).unwrap(),
Provider::Anthropic
);
assert_eq!(
resolve_provider(None, false, true).unwrap(),
Provider::OpenAi
);
}
#[test]
fn resolve_role_defaults_to_memory_extract() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let saved = std::env::var_os("OXIBRAIN_LLM_ROLE");
unsafe {
std::env::remove_var("OXIBRAIN_LLM_ROLE");
}
let got = resolve_role();
unsafe {
if let Some(v) = saved {
std::env::set_var("OXIBRAIN_LLM_ROLE", v);
}
}
assert_eq!(got, ProfileRole::MemoryExtract);
}
#[test]
fn resolve_role_honours_env_when_recognised() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let saved = std::env::var_os("OXIBRAIN_LLM_ROLE");
unsafe {
std::env::set_var("OXIBRAIN_LLM_ROLE", "coding.primary");
}
let got = resolve_role();
unsafe {
match saved {
Some(v) => std::env::set_var("OXIBRAIN_LLM_ROLE", v),
None => std::env::remove_var("OXIBRAIN_LLM_ROLE"),
}
}
assert_eq!(got, ProfileRole::CodingPrimary);
}
#[test]
fn extract_role_entry_is_selected() {
use oxibrain::models::{ModelEntry, ModelRole};
let mk = |role: ModelRole, name: &str| ModelEntry {
role,
name: name.into(),
url: String::new(),
digest: format!("d-{name}"),
size_mb: 1,
license: String::new(),
file: format!("{name}.gguf"),
};
let entries = vec![
mk(ModelRole::Embed, "bge-m3"),
mk(ModelRole::Extract, "qwen2.5-1.5b-instruct"),
];
let got = extract_entry(&entries).expect("extract entry");
assert_eq!(got.name, "qwen2.5-1.5b-instruct");
assert_eq!(got.digest, "d-qwen2.5-1.5b-instruct");
}
}