rust_llm 2.1.1

//! Port of `lib/ruby_llm/configuration.rb`.

use std::collections::HashMap;
use std::sync::{Arc, LazyLock, RwLock};
use std::time::Duration;

/// `RubyLLM::Configuration`. Provider keys live in `values` under RubyLLM's option names
/// (`openai_api_key`, `anthropic_api_base`, ...) so every provider reads its settings the same way.
#[derive(Clone)]
pub struct Config {
    pub default_model: String,
    pub default_embedding_model: String,
    /// `default_image_model`: the model `rust_llm::paint` uses when none is given.
    pub default_image_model: String,
    /// `default_judgment_model`: what `judge` uses; judges never fall back to the chat model.
    pub default_judgment_model: String,
    /// `default_speech_model`: what `rust_llm::speak` uses when no model is given.
    pub default_speech_model: String,
    /// `default_transcription_model`: what `rust_llm::transcribe` uses when no model is given.
    pub default_transcription_model: String,
    /// `default_moderation_model`: what `rust_llm::moderate` uses when no model is given.
    pub default_moderation_model: String,
    /// `default_ocr_model`: what `rust_llm::ocr` uses when no model is given.
    pub default_ocr_model: String,
    /// `default_video_model`: what `rust_llm::animate` uses when no model is given.
    pub default_video_model: String,
    /// `video_generation_timeout`: how long `VideoJob::wait` polls before giving up.
    pub video_generation_timeout: Duration,
    /// `video_generation_poll_interval`: the pause between `VideoJob::wait` polls.
    pub video_generation_poll_interval: Duration,
    pub request_timeout: Duration,
    pub max_retries: u32,
    pub retry_interval: f64,
    pub retry_backoff_factor: f64,
    pub retry_interval_randomness: f64,
    pub retry_max_interval: f64,
    /// `http_proxy`: proxy URL for every request (`Connection.basic` and provider connections).
    pub http_proxy: Option<String>,
    pub tool_concurrency: bool,
    /// `auto_upload_large_files`: upload oversized local attachments to the provider's Files API.
    pub auto_upload_large_files: bool,
    /// `instrumenter`: receives every `*.rust_llm` event (`rust_llm::instrumentation`).
    pub instrumenter: Option<Arc<dyn crate::instrumentation::Instrumenter>>,
    /// `model_registry_file`: where `Models::refresh` saves the registry and where the registry
    /// loads from before falling back to the bundled copy. `None` keeps it in memory only.
    pub model_registry_file: Option<std::path::PathBuf>,
    /// `model_registry_store`: where the registry lives instead of `model_registry_file`, such as
    /// an application's database (`rust_llm::models::registry::ModelRegistryStore`).
    pub model_registry_store: Option<Arc<dyn crate::models::registry::ModelRegistryStore>>,
    /// `Prompt.roots`: the directories `render_prompt` searches after `app/prompts`, in order.
    pub prompt_roots: Vec<std::path::PathBuf>,
    /// `mcp_credential_store`: where MCP OAuth credentials live. `None` keeps them in memory;
    /// `rust_llm_loco::McpCredentialStore` keeps them in `rust_llm_mcp_credentials`.
    pub mcp_credential_store: Option<Arc<dyn crate::mcp::CredentialStore>>,
    /// `mcp_client_name`: the client name given when registering with an MCP server's
    /// authorization server. RubyLLM says `"RubyLLM"`; the port says `"RustLLM"`.
    pub mcp_client_name: String,
    /// `mcp_client_id`: the HTTPS URL of your OAuth client metadata document, used instead of
    /// registering by authorization servers that support client ID metadata documents.
    pub mcp_client_id: Option<String>,
    /// `batch_store`: where submitted chat batches are persisted, so `Batch::find` can hand one
    /// back without contacting the provider. `None` keeps nothing (`crate::batch::BatchStore`).
    pub batch_store: Option<Arc<dyn crate::batch::BatchStore>>,
    /// `Accounting::Usage.ledger`: where attempts no chat records are persisted (one-shot
    /// operations, finished video jobs, chats without a record). `None` keeps them in memory
    /// (`crate::accounting::UsageLedger`).
    pub usage_ledger: Option<Arc<dyn crate::accounting::UsageLedger>>,
    values: HashMap<String, String>,
}

/// `Configuration#inspect`: credentials never reach logs or error output. Like RubyLLM, options
/// ending in `_id`, `_key`, `_secret`, `_token`, or `_credential_provider` are left out entirely.
impl std::fmt::Debug for Config {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        const SECRET: &[&str] = &["_id", "_key", "_secret", "_token", "_credential_provider"];
        let mut values: Vec<(&String, &String)> = self
            .values
            .iter()
            .filter(|(k, _)| !SECRET.iter().any(|s| k.ends_with(s)))
            .collect();
        values.sort();
        f.debug_struct("Config")
            .field("default_model", &self.default_model)
            .field("default_embedding_model", &self.default_embedding_model)
            .field("default_image_model", &self.default_image_model)
            .field("default_judgment_model", &self.default_judgment_model)
            .field("request_timeout", &self.request_timeout)
            .field("max_retries", &self.max_retries)
            .field("http_proxy", &self.http_proxy.as_ref().map(|_| "[set]"))
            .field("tool_concurrency", &self.tool_concurrency)
            .field("auto_upload_large_files", &self.auto_upload_large_files)
            .field("instrumenter", &self.instrumenter.is_some())
            .field("mcp_client_name", &self.mcp_client_name)
            .field("values", &values)
            .finish_non_exhaustive()
    }
}

/// Provider options RubyLLM reads, and the environment variable each falls back to.
const PROVIDER_OPTIONS: &[&str] = &[
    "openai_api_key",
    "openai_api_base",
    "openai_organization_id",
    "openai_project_id",
    "openai_use_system_role",
    "anthropic_api_key",
    "anthropic_api_base",
    "gemini_api_key",
    "gemini_api_base",
    "deepseek_api_key",
    "deepseek_api_base",
    "mistral_api_key",
    "mistral_api_base",
    "openrouter_api_key",
    "openrouter_api_base",
    "openrouter_app_url",
    "openrouter_app_name",
    "xai_api_key",
    "xai_api_base",
    "perplexity_api_key",
    "perplexity_api_base",
    "ollama_api_base",
    "ollama_api_key",
    "ollama_cloud_api_key",
    "ollama_cloud_api_base",
    "gpustack_api_base",
    "gpustack_api_key",
    "hetzner_api_key",
    "hetzner_api_base",
    "typesafe_api_key",
    "typesafe_api_base",
];

impl Default for Config {
    fn default() -> Self {
        Config {
            default_model: "gpt-5.6".into(),
            default_embedding_model: "text-embedding-3-small".into(),
            default_image_model: "gpt-image-2".into(),
            default_judgment_model: "jev-latest".into(),
            default_speech_model: "gpt-4o-mini-tts-2025-12-15".into(),
            default_transcription_model: "gpt-transcribe".into(),
            default_moderation_model: "omni-moderation-latest".into(),
            default_ocr_model: "mistral-ocr-latest".into(),
            default_video_model: "grok-imagine-video-1.5".into(),
            video_generation_timeout: Duration::from_secs(600),
            video_generation_poll_interval: Duration::from_secs(5),
            request_timeout: Duration::from_secs(300),
            max_retries: 3,
            retry_interval: 0.1,
            retry_backoff_factor: 2.0,
            retry_interval_randomness: 0.5,
            retry_max_interval: 30.0,
            http_proxy: None,
            tool_concurrency: false,
            auto_upload_large_files: true,
            instrumenter: None,
            model_registry_file: crate::models::registry::cache_path(),
            model_registry_store: None,
            prompt_roots: Vec::new(),
            mcp_credential_store: None,
            mcp_client_name: "RustLLM".into(),
            mcp_client_id: None,
            batch_store: None,
            usage_ledger: None,
            values: HashMap::new(),
        }
    }
}

impl Config {
    /// Defaults plus any `OPENAI_API_KEY`-style variables present in the environment.
    pub fn from_env() -> Config {
        let mut config = Config::default();
        for option in PROVIDER_OPTIONS {
            if let Ok(value) = std::env::var(option.to_uppercase())
                && !value.is_empty()
            {
                config.values.insert(option.to_string(), value);
            }
        }
        if let Ok(model) = std::env::var("RUST_LLM_DEFAULT_MODEL") {
            config.default_model = model;
        }
        config
    }

    pub fn get(&self, option: &str) -> Option<&str> {
        // `normalize_blank`: a whitespace-only value is unset, like `value.strip.empty?`.
        self.values
            .get(option)
            .map(String::as_str)
            .filter(|v| !v.trim().is_empty())
    }

    pub fn set(&mut self, option: impl Into<String>, value: impl Into<String>) -> &mut Self {
        self.values.insert(option.into(), value.into());
        self
    }

    pub fn openai_api_key(&mut self, v: impl Into<String>) -> &mut Self {
        self.set("openai_api_key", v)
    }
    pub fn anthropic_api_key(&mut self, v: impl Into<String>) -> &mut Self {
        self.set("anthropic_api_key", v)
    }
    pub fn gemini_api_key(&mut self, v: impl Into<String>) -> &mut Self {
        self.set("gemini_api_key", v)
    }
    pub fn deepseek_api_key(&mut self, v: impl Into<String>) -> &mut Self {
        self.set("deepseek_api_key", v)
    }
    pub fn openrouter_api_key(&mut self, v: impl Into<String>) -> &mut Self {
        self.set("openrouter_api_key", v)
    }
    pub fn ollama_api_base(&mut self, v: impl Into<String>) -> &mut Self {
        self.set("ollama_api_base", v)
    }
}

static CONFIG: LazyLock<RwLock<Arc<Config>>> =
    LazyLock::new(|| RwLock::new(Arc::new(Config::from_env())));

/// `RubyLLM.configure { |config| ... }`.
pub fn configure(f: impl FnOnce(&mut Config)) {
    let mut guard = CONFIG.write().unwrap(); // poisoned lock only
    let mut config = (**guard).clone();
    f(&mut config);
    *guard = Arc::new(config);
}

/// `RubyLLM.config`.
pub fn config() -> Arc<Config> {
    CONFIG.read().unwrap().clone() // poisoned lock only
}

/// The configuration, or `None` while `configure` holds it (so loading the model registry from
/// inside a `configure` closure falls back to the bundled copy instead of deadlocking).
pub(crate) fn try_config() -> Option<Arc<Config>> {
    CONFIG.try_read().ok().map(|c| c.clone())
}