use anyhow::Result;
use colored::Colorize;
use crate::config::{Config, LlmConfig};
use crate::exit_codes::*;
#[derive(Debug, Clone)]
pub enum LlmProvider {
OpenAI {
model: String,
api_key: Option<String>,
},
Anthropic {
model: String,
api_key: Option<String>,
},
Ollama { endpoint: String, model: String },
Custom {
endpoint: String,
model: String,
api_key: Option<String>,
},
}
#[derive(Debug)]
pub struct ConfigShowArgs {
pub show_secrets: bool,
}
#[derive(Debug)]
pub enum ConfigLlmArgs {
Set(LlmProvider),
Show { show_secrets: bool },
Remove,
}
pub fn execute_show(args: ConfigShowArgs) -> Result<i32> {
let config = match Config::load() {
Ok(config) => config,
Err(_) => {
eprintln!(
"{} No configuration found. Run `unfault login` first.",
"Error:".red().bold()
);
return Ok(EXIT_CONFIG_ERROR);
}
};
println!();
println!("{}", "Unfault Configuration".bold().underline());
println!();
println!("{}", "LLM (BYOLLM)".cyan().bold());
if let Some(ref llm) = config.llm {
println!(" {} {}", "Provider:".dimmed(), llm.provider);
println!(" {} {}", "Endpoint:".dimmed(), llm.endpoint);
println!(" {} {}", "Model:".dimmed(), llm.model);
if let Some(ref env_var) = llm.api_key_env {
let has_key = std::env::var(env_var).is_ok();
let status = if has_key {
"✓ set".green().to_string()
} else {
"✗ not set".red().to_string()
};
println!(" {} {} ({})", "API Key Env:".dimmed(), env_var, status);
}
if llm.api_key.is_some() {
let display = if args.show_secrets {
llm.api_key.as_ref().unwrap().clone()
} else {
llm.masked_api_key().unwrap_or_else(|| "****".to_string())
};
println!(" {} {}", "API Key:".dimmed(), display);
}
let ready = if llm.is_ready() {
"✓ ready".green()
} else {
"✗ not ready (API key missing)".red()
};
println!(" {} {}", "Status:".dimmed(), ready);
} else {
println!(" {}", "Not configured".dimmed());
println!();
println!(" {} Configure with:", "→".cyan());
println!(" unfault config llm openai --model gpt-4");
println!(" unfault config llm anthropic --model claude-3-5-sonnet-latest");
println!(
" unfault config llm ollama --endpoint http://localhost:11434 --model llama3.2"
);
}
println!();
Ok(EXIT_SUCCESS)
}
pub fn execute_llm(args: ConfigLlmArgs) -> Result<i32> {
match args {
ConfigLlmArgs::Set(provider) => set_llm_config(provider),
ConfigLlmArgs::Show { show_secrets } => show_llm_config(show_secrets),
ConfigLlmArgs::Remove => remove_llm_config(),
}
}
fn set_llm_config(provider: LlmProvider) -> Result<i32> {
let mut config = match Config::load() {
Ok(config) => config,
Err(_) => {
eprintln!(
"{} No configuration found. Run `unfault login` first.",
"Error:".red().bold()
);
return Ok(EXIT_CONFIG_ERROR);
}
};
let llm_config = match provider {
LlmProvider::OpenAI { model, api_key } => {
let mut cfg = LlmConfig::openai(&model);
if let Some(key) = api_key {
cfg.api_key = Some(key);
}
cfg
}
LlmProvider::Anthropic { model, api_key } => {
let mut cfg = LlmConfig::anthropic(&model);
if let Some(key) = api_key {
cfg.api_key = Some(key);
}
cfg
}
LlmProvider::Ollama { endpoint, model } => LlmConfig::ollama(&endpoint, &model),
LlmProvider::Custom {
endpoint,
model,
api_key,
} => {
let mut cfg = LlmConfig::custom(&endpoint, &model);
cfg.api_key = api_key;
cfg
}
};
let provider_name = llm_config.provider.clone();
let model_name = llm_config.model.clone();
config.llm = Some(llm_config.clone());
config.save()?;
println!();
println!("{} LLM configured successfully!", "✓".green().bold());
println!();
println!(" {} {}", "Provider:".dimmed(), provider_name);
println!(" {} {}", "Model:".dimmed(), model_name);
println!(" {} {}", "Endpoint:".dimmed(), llm_config.endpoint);
if !llm_config.is_ready() {
println!();
eprintln!(
"{} API key not found. Set the {} environment variable.",
"⚠".yellow().bold(),
llm_config.api_key_env.as_deref().unwrap_or("API_KEY")
);
} else {
println!();
println!(" {} Ready to use with `unfault ask`", "→".cyan());
}
println!();
Ok(EXIT_SUCCESS)
}
fn show_llm_config(show_secrets: bool) -> Result<i32> {
let config = match Config::load() {
Ok(config) => config,
Err(_) => {
eprintln!(
"{} No configuration found. Run `unfault login` first.",
"Error:".red().bold()
);
return Ok(EXIT_CONFIG_ERROR);
}
};
println!();
println!("{}", "LLM Configuration".bold().underline());
println!();
if let Some(ref llm) = config.llm {
println!(" {} {}", "Provider:".dimmed(), llm.provider);
println!(" {} {}", "Endpoint:".dimmed(), llm.endpoint);
println!(" {} {}", "Model:".dimmed(), llm.model);
if let Some(ref env_var) = llm.api_key_env {
let has_key = std::env::var(env_var).is_ok();
let status = if has_key {
"✓ set".green().to_string()
} else {
"✗ not set".red().to_string()
};
println!(" {} {} ({})", "API Key Env:".dimmed(), env_var, status);
}
if llm.api_key.is_some() {
let display = if show_secrets {
llm.api_key.as_ref().unwrap().clone()
} else {
llm.masked_api_key().unwrap_or_else(|| "****".to_string())
};
println!(" {} {}", "API Key:".dimmed(), display);
}
let ready = if llm.is_ready() {
"✓ ready".green()
} else {
"✗ not ready (API key missing)".red()
};
println!();
println!(" {} {}", "Status:".dimmed(), ready);
} else {
println!(" {}", "Not configured".dimmed());
println!();
println!(" {} Configure with:", "→".cyan());
println!(" unfault config llm openai --model gpt-4");
println!(" unfault config llm anthropic --model claude-3-5-sonnet-latest");
println!(
" unfault config llm ollama --endpoint http://localhost:11434 --model llama3.2"
);
}
println!();
Ok(EXIT_SUCCESS)
}
fn remove_llm_config() -> Result<i32> {
let mut config = match Config::load() {
Ok(config) => config,
Err(_) => {
eprintln!(
"{} No configuration found. Run `unfault login` first.",
"Error:".red().bold()
);
return Ok(EXIT_CONFIG_ERROR);
}
};
if config.llm.is_none() {
println!();
println!("{} LLM configuration is not set.", "ℹ".blue());
println!();
return Ok(EXIT_SUCCESS);
}
config.remove_llm();
config.save()?;
println!();
println!("{} LLM configuration removed.", "✓".green().bold());
println!();
Ok(EXIT_SUCCESS)
}
#[cfg(test)]
fn mask_key(key: &str) -> String {
if key.len() > 8 {
format!("{}...{}", &key[..4], &key[key.len() - 4..])
} else {
"****".to_string()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_mask_key_long() {
let key = "sk_live_1234567890abcdef";
let masked = mask_key(key);
assert_eq!(masked, "sk_l...cdef");
}
#[test]
fn test_mask_key_short() {
let key = "short";
let masked = mask_key(key);
assert_eq!(masked, "****");
}
#[test]
fn test_llm_provider_openai() {
let provider = LlmProvider::OpenAI {
model: "gpt-4".to_string(),
api_key: None,
};
match provider {
LlmProvider::OpenAI { model, api_key } => {
assert_eq!(model, "gpt-4");
assert!(api_key.is_none());
}
_ => panic!("Expected OpenAI provider"),
}
}
#[test]
fn test_llm_provider_anthropic() {
let provider = LlmProvider::Anthropic {
model: "claude-3-5-sonnet-latest".to_string(),
api_key: Some("test-key".to_string()),
};
match provider {
LlmProvider::Anthropic { model, api_key } => {
assert_eq!(model, "claude-3-5-sonnet-latest");
assert_eq!(api_key, Some("test-key".to_string()));
}
_ => panic!("Expected Anthropic provider"),
}
}
#[test]
fn test_llm_provider_ollama() {
let provider = LlmProvider::Ollama {
endpoint: "http://localhost:11434".to_string(),
model: "llama3.2".to_string(),
};
match provider {
LlmProvider::Ollama { endpoint, model } => {
assert_eq!(endpoint, "http://localhost:11434");
assert_eq!(model, "llama3.2");
}
_ => panic!("Expected Ollama provider"),
}
}
#[test]
fn test_llm_provider_custom() {
let provider = LlmProvider::Custom {
endpoint: "https://api.example.com/v1".to_string(),
model: "custom-model".to_string(),
api_key: None,
};
match provider {
LlmProvider::Custom {
endpoint,
model,
api_key,
} => {
assert_eq!(endpoint, "https://api.example.com/v1");
assert_eq!(model, "custom-model");
assert!(api_key.is_none());
}
_ => panic!("Expected Custom provider"),
}
}
}