feat(cli): add provider management and model listing commands and integrate them into the CLI
This commit is contained in:
479
crates/owlen-cli/src/commands/cloud.rs
Normal file
479
crates/owlen-cli/src/commands/cloud.rs
Normal file
@@ -0,0 +1,479 @@
|
||||
use std::ffi::OsStr;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
use clap::Subcommand;
|
||||
use owlen_core::LlmProvider;
|
||||
use owlen_core::ProviderConfig;
|
||||
use owlen_core::config::{
|
||||
self as core_config, Config, OLLAMA_CLOUD_API_KEY_ENV, OLLAMA_CLOUD_BASE_URL,
|
||||
OLLAMA_CLOUD_ENDPOINT_KEY, OLLAMA_MODE_KEY,
|
||||
};
|
||||
use owlen_core::credentials::{ApiCredentials, CredentialManager, OLLAMA_CLOUD_CREDENTIAL_ID};
|
||||
use owlen_core::encryption;
|
||||
use owlen_core::providers::OllamaProvider;
|
||||
use owlen_core::storage::StorageManager;
|
||||
use serde_json::Value;
|
||||
|
||||
const DEFAULT_CLOUD_ENDPOINT: &str = OLLAMA_CLOUD_BASE_URL;
|
||||
const CLOUD_ENDPOINT_KEY: &str = OLLAMA_CLOUD_ENDPOINT_KEY;
|
||||
const CLOUD_PROVIDER_KEY: &str = "ollama_cloud";
|
||||
|
||||
#[derive(Debug, Subcommand)]
|
||||
pub enum CloudCommand {
|
||||
/// Configure Ollama Cloud credentials
|
||||
Setup {
|
||||
/// API key passed directly on the command line (prompted when omitted)
|
||||
#[arg(long)]
|
||||
api_key: Option<String>,
|
||||
/// Override the cloud endpoint (default: https://ollama.com)
|
||||
#[arg(long)]
|
||||
endpoint: Option<String>,
|
||||
/// Provider name to configure (default: ollama_cloud)
|
||||
#[arg(long, default_value = "ollama_cloud")]
|
||||
provider: String,
|
||||
/// Overwrite the provider base URL with the cloud endpoint
|
||||
#[arg(long)]
|
||||
force_cloud_base_url: bool,
|
||||
},
|
||||
/// Check connectivity to Ollama Cloud
|
||||
Status {
|
||||
/// Provider name to check (default: ollama_cloud)
|
||||
#[arg(long, default_value = "ollama_cloud")]
|
||||
provider: String,
|
||||
},
|
||||
/// List available cloud-hosted models
|
||||
Models {
|
||||
/// Provider name to query (default: ollama_cloud)
|
||||
#[arg(long, default_value = "ollama_cloud")]
|
||||
provider: String,
|
||||
},
|
||||
/// Remove stored Ollama Cloud credentials
|
||||
Logout {
|
||||
/// Provider name to clear (default: ollama_cloud)
|
||||
#[arg(long, default_value = "ollama_cloud")]
|
||||
provider: String,
|
||||
},
|
||||
}
|
||||
|
||||
pub async fn run_cloud_command(command: CloudCommand) -> Result<()> {
|
||||
match command {
|
||||
CloudCommand::Setup {
|
||||
api_key,
|
||||
endpoint,
|
||||
provider,
|
||||
force_cloud_base_url,
|
||||
} => setup(provider, api_key, endpoint, force_cloud_base_url).await,
|
||||
CloudCommand::Status { provider } => status(provider).await,
|
||||
CloudCommand::Models { provider } => models(provider).await,
|
||||
CloudCommand::Logout { provider } => logout(provider).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn setup(
|
||||
provider: String,
|
||||
api_key: Option<String>,
|
||||
endpoint: Option<String>,
|
||||
force_cloud_base_url: bool,
|
||||
) -> Result<()> {
|
||||
let provider = canonical_provider_name(&provider);
|
||||
let mut config = crate::config::try_load_config().unwrap_or_default();
|
||||
let endpoint =
|
||||
normalize_endpoint(&endpoint.unwrap_or_else(|| DEFAULT_CLOUD_ENDPOINT.to_string()));
|
||||
|
||||
let base_changed = {
|
||||
let entry = ensure_provider_entry(&mut config, &provider);
|
||||
entry.enabled = true;
|
||||
configure_cloud_endpoint(entry, &endpoint, force_cloud_base_url)
|
||||
};
|
||||
|
||||
let key = match api_key {
|
||||
Some(value) if !value.trim().is_empty() => value,
|
||||
_ => {
|
||||
let prompt = format!("Enter API key for {provider}: ");
|
||||
encryption::prompt_password(&prompt)?
|
||||
}
|
||||
};
|
||||
|
||||
if config.privacy.encrypt_local_data {
|
||||
let storage = Arc::new(StorageManager::new().await?);
|
||||
let manager = unlock_credential_manager(&config, storage.clone())?;
|
||||
let credentials = ApiCredentials {
|
||||
api_key: key.clone(),
|
||||
endpoint: endpoint.clone(),
|
||||
};
|
||||
manager
|
||||
.store_credentials(OLLAMA_CLOUD_CREDENTIAL_ID, &credentials)
|
||||
.await?;
|
||||
// Ensure plaintext key is not persisted to disk.
|
||||
if let Some(entry) = config.providers.get_mut(&provider) {
|
||||
entry.api_key = None;
|
||||
}
|
||||
} else if let Some(entry) = config.providers.get_mut(&provider) {
|
||||
entry.api_key = Some(key.clone());
|
||||
}
|
||||
|
||||
crate::config::save_config(&config)?;
|
||||
println!("Saved Ollama configuration for provider '{provider}'.");
|
||||
if config.privacy.encrypt_local_data {
|
||||
println!("API key stored securely in the encrypted credential vault.");
|
||||
} else {
|
||||
println!("API key stored in plaintext configuration (encryption disabled).");
|
||||
}
|
||||
if !force_cloud_base_url && !base_changed {
|
||||
println!(
|
||||
"Local base URL preserved; cloud endpoint stored as {}.",
|
||||
CLOUD_ENDPOINT_KEY
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn status(provider: String) -> Result<()> {
|
||||
let provider = canonical_provider_name(&provider);
|
||||
let mut config = crate::config::try_load_config().unwrap_or_default();
|
||||
let storage = Arc::new(StorageManager::new().await?);
|
||||
let manager = if config.privacy.encrypt_local_data {
|
||||
Some(unlock_credential_manager(&config, storage.clone())?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let api_key = hydrate_api_key(&mut config, manager.as_ref()).await?;
|
||||
{
|
||||
let entry = ensure_provider_entry(&mut config, &provider);
|
||||
entry.enabled = true;
|
||||
configure_cloud_endpoint(entry, DEFAULT_CLOUD_ENDPOINT, false);
|
||||
}
|
||||
|
||||
let provider_cfg = config
|
||||
.provider(&provider)
|
||||
.cloned()
|
||||
.ok_or_else(|| anyhow!("Provider '{provider}' is not configured"))?;
|
||||
|
||||
let endpoint =
|
||||
resolve_cloud_endpoint(&provider_cfg).unwrap_or_else(|| DEFAULT_CLOUD_ENDPOINT.to_string());
|
||||
let mut runtime_cfg = provider_cfg.clone();
|
||||
runtime_cfg.base_url = Some(endpoint.clone());
|
||||
runtime_cfg.extra.insert(
|
||||
OLLAMA_MODE_KEY.to_string(),
|
||||
Value::String("cloud".to_string()),
|
||||
);
|
||||
|
||||
let ollama = OllamaProvider::from_config(&runtime_cfg, Some(&config.general))
|
||||
.with_context(|| "Failed to construct Ollama provider. Run `owlen cloud setup` first.")?;
|
||||
|
||||
match ollama.health_check().await {
|
||||
Ok(_) => {
|
||||
println!("✓ Connected to {provider} ({})", endpoint);
|
||||
if api_key.is_none() && config.privacy.encrypt_local_data {
|
||||
println!(
|
||||
"Warning: No API key stored; connection succeeded via environment variables."
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
println!("✗ Failed to reach {provider}: {err}");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn models(provider: String) -> Result<()> {
|
||||
let provider = canonical_provider_name(&provider);
|
||||
let mut config = crate::config::try_load_config().unwrap_or_default();
|
||||
let storage = Arc::new(StorageManager::new().await?);
|
||||
let manager = if config.privacy.encrypt_local_data {
|
||||
Some(unlock_credential_manager(&config, storage.clone())?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
hydrate_api_key(&mut config, manager.as_ref()).await?;
|
||||
|
||||
{
|
||||
let entry = ensure_provider_entry(&mut config, &provider);
|
||||
entry.enabled = true;
|
||||
configure_cloud_endpoint(entry, DEFAULT_CLOUD_ENDPOINT, false);
|
||||
}
|
||||
|
||||
let provider_cfg = config
|
||||
.provider(&provider)
|
||||
.cloned()
|
||||
.ok_or_else(|| anyhow!("Provider '{provider}' is not configured"))?;
|
||||
|
||||
let endpoint =
|
||||
resolve_cloud_endpoint(&provider_cfg).unwrap_or_else(|| DEFAULT_CLOUD_ENDPOINT.to_string());
|
||||
let mut runtime_cfg = provider_cfg.clone();
|
||||
runtime_cfg.base_url = Some(endpoint);
|
||||
runtime_cfg.extra.insert(
|
||||
OLLAMA_MODE_KEY.to_string(),
|
||||
Value::String("cloud".to_string()),
|
||||
);
|
||||
|
||||
let ollama = OllamaProvider::from_config(&runtime_cfg, Some(&config.general))
|
||||
.with_context(|| "Failed to construct Ollama provider. Run `owlen cloud setup` first.")?;
|
||||
|
||||
match ollama.list_models().await {
|
||||
Ok(models) => {
|
||||
if models.is_empty() {
|
||||
println!("No cloud models reported by '{}'.", provider);
|
||||
} else {
|
||||
println!("Models available via '{}':", provider);
|
||||
for model in models {
|
||||
if let Some(description) = &model.description {
|
||||
println!(" - {} ({})", model.id, description);
|
||||
} else {
|
||||
println!(" - {}", model.id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
bail!("Failed to list models: {err}");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn logout(provider: String) -> Result<()> {
|
||||
let provider = canonical_provider_name(&provider);
|
||||
let mut config = crate::config::try_load_config().unwrap_or_default();
|
||||
let storage = Arc::new(StorageManager::new().await?);
|
||||
|
||||
if config.privacy.encrypt_local_data {
|
||||
let manager = unlock_credential_manager(&config, storage.clone())?;
|
||||
manager
|
||||
.delete_credentials(OLLAMA_CLOUD_CREDENTIAL_ID)
|
||||
.await?;
|
||||
}
|
||||
|
||||
if let Some(entry) = config.providers.get_mut(&provider) {
|
||||
entry.api_key = None;
|
||||
entry.enabled = false;
|
||||
}
|
||||
|
||||
crate::config::save_config(&config)?;
|
||||
println!("Cleared credentials for provider '{provider}'.");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ensure_provider_entry<'a>(config: &'a mut Config, provider: &str) -> &'a mut ProviderConfig {
|
||||
core_config::ensure_provider_config_mut(config, provider)
|
||||
}
|
||||
|
||||
fn configure_cloud_endpoint(entry: &mut ProviderConfig, endpoint: &str, force: bool) -> bool {
|
||||
let normalized = normalize_endpoint(endpoint);
|
||||
let previous_base = entry.base_url.clone();
|
||||
entry.extra.insert(
|
||||
CLOUD_ENDPOINT_KEY.to_string(),
|
||||
Value::String(normalized.clone()),
|
||||
);
|
||||
|
||||
if entry.api_key_env.is_none() {
|
||||
entry.api_key_env = Some(OLLAMA_CLOUD_API_KEY_ENV.to_string());
|
||||
}
|
||||
|
||||
if force
|
||||
|| entry
|
||||
.base_url
|
||||
.as_ref()
|
||||
.map(|value| value.trim().is_empty())
|
||||
.unwrap_or(true)
|
||||
{
|
||||
entry.base_url = Some(normalized.clone());
|
||||
}
|
||||
|
||||
if force {
|
||||
entry.enabled = true;
|
||||
}
|
||||
|
||||
entry.base_url != previous_base
|
||||
}
|
||||
|
||||
fn resolve_cloud_endpoint(cfg: &ProviderConfig) -> Option<String> {
|
||||
if let Some(value) = cfg
|
||||
.extra
|
||||
.get(CLOUD_ENDPOINT_KEY)
|
||||
.and_then(|value| value.as_str())
|
||||
.map(normalize_endpoint)
|
||||
{
|
||||
return Some(value);
|
||||
}
|
||||
|
||||
cfg.base_url
|
||||
.as_ref()
|
||||
.map(|value| value.trim_end_matches('/').to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn normalize_endpoint(endpoint: &str) -> String {
|
||||
let trimmed = endpoint.trim().trim_end_matches('/');
|
||||
if trimmed.is_empty() {
|
||||
DEFAULT_CLOUD_ENDPOINT.to_string()
|
||||
} else {
|
||||
trimmed.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn canonical_provider_name(provider: &str) -> String {
|
||||
let normalized = provider.trim().to_ascii_lowercase().replace('-', "_");
|
||||
match normalized.as_str() {
|
||||
"" => CLOUD_PROVIDER_KEY.to_string(),
|
||||
"ollama" => CLOUD_PROVIDER_KEY.to_string(),
|
||||
"ollama_cloud" => CLOUD_PROVIDER_KEY.to_string(),
|
||||
value => value.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn set_env_var<K, V>(key: K, value: V)
|
||||
where
|
||||
K: AsRef<OsStr>,
|
||||
V: AsRef<OsStr>,
|
||||
{
|
||||
// Safety: the CLI updates process-wide environment variables during startup while no
|
||||
// other threads are mutating the environment.
|
||||
unsafe {
|
||||
std::env::set_var(key, value);
|
||||
}
|
||||
}
|
||||
|
||||
fn set_env_if_missing(var: &str, value: &str) {
|
||||
if std::env::var(var)
|
||||
.map(|v| v.trim().is_empty())
|
||||
.unwrap_or(true)
|
||||
{
|
||||
set_env_var(var, value);
|
||||
}
|
||||
}
|
||||
|
||||
fn unlock_credential_manager(
|
||||
config: &Config,
|
||||
storage: Arc<StorageManager>,
|
||||
) -> Result<Arc<CredentialManager>> {
|
||||
if !config.privacy.encrypt_local_data {
|
||||
bail!("Credential manager requested but encryption is disabled");
|
||||
}
|
||||
|
||||
let secure_path = vault_path(&storage)?;
|
||||
let handle = unlock_vault(&secure_path)?;
|
||||
let master_key = Arc::new(handle.data.master_key.clone());
|
||||
Ok(Arc::new(CredentialManager::new(
|
||||
storage,
|
||||
master_key.clone(),
|
||||
)))
|
||||
}
|
||||
|
||||
fn vault_path(storage: &StorageManager) -> Result<PathBuf> {
|
||||
let base_dir = storage
|
||||
.database_path()
|
||||
.parent()
|
||||
.map(|p| p.to_path_buf())
|
||||
.or_else(dirs::data_local_dir)
|
||||
.unwrap_or_else(|| PathBuf::from("."));
|
||||
Ok(base_dir.join("encrypted_data.json"))
|
||||
}
|
||||
|
||||
fn unlock_vault(path: &Path) -> Result<encryption::VaultHandle> {
|
||||
use std::env;
|
||||
|
||||
if path.exists() {
|
||||
if let Some(password) = env::var("OWLEN_MASTER_PASSWORD")
|
||||
.ok()
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|password| !password.is_empty())
|
||||
{
|
||||
return encryption::unlock_with_password(path.to_path_buf(), &password)
|
||||
.context("Failed to unlock vault with OWLEN_MASTER_PASSWORD");
|
||||
}
|
||||
|
||||
for attempt in 0..3 {
|
||||
let password = encryption::prompt_password("Enter master password: ")?;
|
||||
match encryption::unlock_with_password(path.to_path_buf(), &password) {
|
||||
Ok(handle) => {
|
||||
set_env_var("OWLEN_MASTER_PASSWORD", password);
|
||||
return Ok(handle);
|
||||
}
|
||||
Err(err) => {
|
||||
eprintln!("Failed to unlock vault: {err}");
|
||||
if attempt == 2 {
|
||||
return Err(err);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
bail!("Unable to unlock encrypted credential vault");
|
||||
}
|
||||
|
||||
let handle = encryption::unlock_interactive(path.to_path_buf())?;
|
||||
if env::var("OWLEN_MASTER_PASSWORD")
|
||||
.map(|v| v.trim().is_empty())
|
||||
.unwrap_or(true)
|
||||
{
|
||||
let password = encryption::prompt_password("Cache master password for this session: ")?;
|
||||
set_env_var("OWLEN_MASTER_PASSWORD", password);
|
||||
}
|
||||
Ok(handle)
|
||||
}
|
||||
|
||||
async fn hydrate_api_key(
|
||||
config: &mut Config,
|
||||
manager: Option<&Arc<CredentialManager>>,
|
||||
) -> Result<Option<String>> {
|
||||
let credentials = match manager {
|
||||
Some(manager) => manager.get_credentials(OLLAMA_CLOUD_CREDENTIAL_ID).await?,
|
||||
None => None,
|
||||
};
|
||||
|
||||
if let Some(credentials) = credentials {
|
||||
let key = credentials.api_key.trim().to_string();
|
||||
if !key.is_empty() {
|
||||
set_env_if_missing("OLLAMA_API_KEY", &key);
|
||||
set_env_if_missing("OLLAMA_CLOUD_API_KEY", &key);
|
||||
}
|
||||
|
||||
let cfg = core_config::ensure_provider_config_mut(config, CLOUD_PROVIDER_KEY);
|
||||
configure_cloud_endpoint(cfg, &credentials.endpoint, false);
|
||||
return Ok(Some(key));
|
||||
}
|
||||
|
||||
if let Some(key) = config
|
||||
.provider(CLOUD_PROVIDER_KEY)
|
||||
.and_then(|cfg| cfg.api_key.as_ref())
|
||||
.map(|value| value.trim())
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
set_env_if_missing("OLLAMA_API_KEY", key);
|
||||
set_env_if_missing("OLLAMA_CLOUD_API_KEY", key);
|
||||
return Ok(Some(key.to_string()));
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub async fn load_runtime_credentials(
|
||||
config: &mut Config,
|
||||
storage: Arc<StorageManager>,
|
||||
) -> Result<()> {
|
||||
if config.privacy.encrypt_local_data {
|
||||
let manager = unlock_credential_manager(config, storage.clone())?;
|
||||
hydrate_api_key(config, Some(&manager)).await?;
|
||||
} else {
|
||||
hydrate_api_key(config, None).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn canonicalises_provider_names() {
|
||||
assert_eq!(canonical_provider_name("OLLAMA_CLOUD"), CLOUD_PROVIDER_KEY);
|
||||
assert_eq!(canonical_provider_name(" ollama-cloud"), CLOUD_PROVIDER_KEY);
|
||||
assert_eq!(canonical_provider_name(""), CLOUD_PROVIDER_KEY);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user