feat: add provider health checks and improve rig adapters

- load per-provider credentials from YAML profiles and expose a health subcommand that primes env overrides before probing
- enhance rig adapter parsing (code fences, fallback verdicts, Gemini JSON MIME) so health checks survive empty or non-textual responses
- document the workflow, ship an example llm_providers.yaml, and add integration/unit tests for the new paths

Author: HendrikReh

Cargo.toml

@@ -16,6 +16,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 serde_yaml = "0.9"
+json5 = "0.4"
 aho-corasick = "1"
 regex = "1"
 clap = { version = "4", features = ["derive"] }

README.md

@@ -1,16 +1,18 @@
 use std::collections::HashMap;
+use std::env;
 use std::fs as stdfs;
 use std::path::{Path, PathBuf};
 use std::process;
 use std::sync::Arc;
 use std::time::Duration;
 
-use anyhow::{anyhow, Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use clap::{Parser, Subcommand};
 use config::Config;
 use llm_guard_core::{
     build_client, render_report, DefaultScanner, FileRuleRepository, LlmClient, LlmSettings,
-    OutputFormat, RiskBand, RuleKind, RuleRepository, Scanner,
+    OutputFormat, RiskBand, RiskThresholds, RuleKind, RuleRepository, ScanReport, Scanner,
+    ScoreBreakdown,
 };
 use serde::Deserialize;
 use serde_yaml;
@@ -52,6 +54,10 @@ struct Cli {
     )]
     providers_config: PathBuf,
 
+    /// Enable verbose diagnostics (including raw provider payloads on errors).
+    #[arg(long, global = true)]
+    debug: bool,
+
     #[command(subcommand)]
     command: Option<Commands>,
 }
@@ -97,6 +103,15 @@ enum Commands {
         #[arg(long)]
         workspace: Option<String>,
     },
+    /// Execute health checks against configured LLM providers.
+    Health {
+        /// Limit the health check to a single provider name.
+        #[arg(long)]
+        provider: Option<String>,
+        /// Skip the live LLM call; only validate configuration/build steps.
+        #[arg(long)]
+        dry_run: bool,
+    },
 }
 
 #[derive(Debug, Deserialize, Clone)]
@@ -165,8 +180,8 @@ impl ProviderProfiles {
     }
 
     fn prime_env(&self, provider: &str) {
-        if let Some(profile) = self.entries.get(&provider.to_ascii_lowercase()) {
-            maybe_set_env("LLM_GUARD_PROVIDER", Some(provider.to_string()));
+        if let Some(profile) = self.get(provider) {
+            maybe_set_env("LLM_GUARD_PROVIDER", Some(profile.name.clone()));
             maybe_set_env("LLM_GUARD_API_KEY", profile.api_key.clone());
             maybe_set_env("LLM_GUARD_ENDPOINT", profile.endpoint.clone());
             maybe_set_env("LLM_GUARD_MODEL", profile.model.clone());
@@ -186,7 +201,7 @@ impl ProviderProfiles {
     }
 
     fn apply_defaults(&self, provider: &str, settings: &mut LlmSettings) {
-        if let Some(profile) = self.entries.get(&provider.to_ascii_lowercase()) {
+        if let Some(profile) = self.get(provider) {
             if settings.model.is_none() {
                 settings.model = profile.model.clone();
             }
@@ -212,6 +227,58 @@ impl ProviderProfiles {
             }
         }
     }
+
+    fn get(&self, provider: &str) -> Option<&ProviderProfile> {
+        self.entries.get(&provider.to_ascii_lowercase())
+    }
+
+    fn names(&self) -> Vec<String> {
+        self.entries
+            .values()
+            .map(|profile| profile.name.clone())
+            .collect()
+    }
+
+    fn is_empty(&self) -> bool {
+        self.entries.is_empty()
+    }
+}
+
+struct EnvGuard {
+    snapshot: Vec<(String, Option<String>)>,
+}
+
+impl EnvGuard {
+    fn new() -> Self {
+        Self {
+            snapshot: Vec::new(),
+        }
+    }
+
+    fn set(&mut self, key: &str, value: &str) {
+        if !self.snapshot.iter().any(|(k, _)| k == key) {
+            self.snapshot.push((key.to_string(), env::var(key).ok()));
+        }
+        env::set_var(key, value);
+    }
+
+    fn maybe_set(&mut self, key: &str, value: Option<&str>) {
+        if let Some(val) = value {
+            self.set(key, val);
+        }
+    }
+}
+
+impl Drop for EnvGuard {
+    fn drop(&mut self) {
+        for (key, previous) in self.snapshot.drain(..).rev() {
+            if let Some(value) = previous {
+                env::set_var(&key, value);
+            } else {
+                env::remove_var(&key);
+            }
+        }
+    }
 }
 
 #[cfg(test)]
@@ -338,6 +405,11 @@ async fn main() {
 async fn run() -> Result<i32> {
     init_tracing();
     let cli = Cli::parse();
+    if cli.debug {
+        env::set_var("LLM_GUARD_DEBUG", "1");
+    } else {
+        env::remove_var("LLM_GUARD_DEBUG");
+    }
     let provider_profiles = ProviderProfiles::load(&cli.providers_config)?;
     match cli.command.unwrap_or(Commands::ListRules { json: false }) {
         Commands::ListRules { json } => {
@@ -373,6 +445,9 @@ async fn run() -> Result<i32> {
             )
             .await
         }
+        Commands::Health { provider, dry_run } => {
+            run_health(&provider_profiles, provider.as_deref(), !dry_run).await
+        }
     }
 }
 
@@ -642,6 +717,99 @@ fn exit_code_for_band(band: RiskBand) -> i32 {
     }
 }
 
+async fn run_health(
+    profiles: &ProviderProfiles,
+    provider_filter: Option<&str>,
+    perform_call: bool,
+) -> Result<i32> {
+    let mut targets = if let Some(filter) = provider_filter {
+        if let Some(profile) = profiles.get(filter) {
+            vec![profile.name.clone()]
+        } else {
+            vec![filter.to_string()]
+        }
+    } else if !profiles.is_empty() {
+        profiles.names()
+    } else if let Ok(env_provider) = env::var("LLM_GUARD_PROVIDER") {
+        vec![env_provider]
+    } else {
+        bail!("no providers configured; supply --provider or create llm_providers.yaml");
+    };
+
+    targets.sort();
+    targets.dedup();
+
+    let mut failed = false;
+    for provider in targets {
+        println!("Checking provider {provider}...");
+        match check_provider(profiles, &provider, perform_call).await {
+            Ok(()) => println!("  ok"),
+            Err(err) => {
+                failed = true;
+                eprintln!("  failed: {err:#}");
+            }
+        }
+    }
+
+    Ok(if failed { 1 } else { 0 })
+}
+
+async fn check_provider(
+    profiles: &ProviderProfiles,
+    provider: &str,
+    perform_call: bool,
+) -> Result<()> {
+    let profile_snapshot = profiles.get(provider).cloned();
+    let canonical_provider = profile_snapshot
+        .as_ref()
+        .map(|p| p.name.clone())
+        .unwrap_or_else(|| provider.to_string());
+
+    let mut guard = EnvGuard::new();
+    guard.set("LLM_GUARD_PROVIDER", &canonical_provider);
+    if let Some(profile) = profile_snapshot.as_ref() {
+        guard.maybe_set("LLM_GUARD_API_KEY", profile.api_key.as_deref());
+        guard.maybe_set("LLM_GUARD_ENDPOINT", profile.endpoint.as_deref());
+        guard.maybe_set("LLM_GUARD_MODEL", profile.model.as_deref());
+        guard.maybe_set("LLM_GUARD_DEPLOYMENT", profile.deployment.as_deref());
+        guard.maybe_set("LLM_GUARD_PROJECT", profile.project.as_deref());
+        guard.maybe_set("LLM_GUARD_WORKSPACE", profile.workspace.as_deref());
+        if let Some(timeout) = profile.timeout_secs {
+            guard.set("LLM_GUARD_TIMEOUT_SECS", &timeout.to_string());
+        }
+        if let Some(retries) = profile.max_retries {
+            guard.set("LLM_GUARD_MAX_RETRIES", &retries.to_string());
+        }
+        guard.maybe_set("LLM_GUARD_API_VERSION", profile.api_version.as_deref());
+    }
+
+    let mut settings = LlmSettings::from_env()?;
+    let provider_for_defaults = settings.provider.clone();
+    profiles.apply_defaults(&provider_for_defaults, &mut settings);
+    drop(guard);
+
+    let client = build_client(&settings)?;
+    if perform_call {
+        let report = dummy_report();
+        let _ = client
+            .enrich("Health check probe", &report)
+            .await
+            .context("LLM enrich call failed")?;
+    }
+
+    Ok(())
+}
+
+fn dummy_report() -> ScanReport {
+    ScanReport::from_breakdown(
+        Vec::new(),
+        0,
+        None,
+        ScoreBreakdown::default(),
+        &RiskThresholds::default(),
+    )
+}
+
 fn init_tracing() {
     let env_filter =
         EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info,tokio=warn"));

crates/llm-guard-cli/src/main.rs

@@ -0,0 +1,44 @@
+use assert_cmd::Command;
+use once_cell::sync::Lazy;
+use predicates::str::contains;
+use std::env;
+use std::fs::write;
+use std::sync::Mutex;
+use tempfile;
+
+static ENV_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
+
+fn reset_env() {
+    env::remove_var("LLM_GUARD_PROVIDER");
+    env::remove_var("LLM_GUARD_API_KEY");
+    env::remove_var("LLM_GUARD_ENDPOINT");
+    env::remove_var("LLM_GUARD_MODEL");
+    env::remove_var("LLM_GUARD_DEPLOYMENT");
+    env::remove_var("LLM_GUARD_PROJECT");
+    env::remove_var("LLM_GUARD_WORKSPACE");
+    env::remove_var("LLM_GUARD_TIMEOUT_SECS");
+    env::remove_var("LLM_GUARD_MAX_RETRIES");
+    env::remove_var("LLM_GUARD_API_VERSION");
+    env::remove_var("LLM_GUARD_DEBUG");
+}
+
+#[test]
+fn health_check_with_noop_profile() {
+    let _guard = ENV_LOCK.lock().unwrap();
+    reset_env();
+
+    let file = tempfile::Builder::new().suffix(".yaml").tempfile().unwrap();
+
+    write(file.path(), "providers:\n  - name: \"noop\"\n").unwrap();
+
+    let mut cmd = Command::cargo_bin("llm-guard-cli").unwrap();
+    cmd.args([
+        "--providers-config",
+        file.path().to_str().unwrap(),
+        "health",
+    ])
+    .assert()
+    .success()
+    .stdout(contains("Checking provider noop"))
+    .stdout(contains("ok"));
+}

crates/llm-guard-cli/tests/health.rs

@@ -18,6 +18,7 @@ tracing.workspace = true
 reqwest.workspace = true
 tokio.workspace = true
 rig-core = "0.22.0"
+json5.workspace = true
 
 [dev-dependencies]
 tempfile = "3"