ForgeLM/crates/localllm-cli/src/main.rs

use clap::{Args, Parser, Subcommand, ValueEnum};

#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
enum OutputFormat {
    Text,
    Json,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
enum AuthMode {
    Off,
    LocalOnly,
    Protected,
}

#[derive(Debug, Parser)]
#[command(name = "localllm", about = "Manage the local LLM backend")]
struct Cli {
    #[arg(long, global = true, default_value = "http://127.0.0.1:8080")]
    server: String,

    #[arg(long, global = true)]
    api_key: Option<String>,

    #[arg(long, global = true, default_value = "./models")]
    models_dir: String,

    #[arg(long, global = true, value_enum, default_value_t = OutputFormat::Text)]
    format: OutputFormat,

    #[command(subcommand)]
    command: Commands,
}

#[derive(Debug, Subcommand)]
enum Commands {
    Status,
    Models(ModelsCommand),
    Auth(AuthCommand),
    Config(ConfigCommand),
    Bench(BenchCommand),
    Logs(LogsArgs),
    Doctor,
}

#[derive(Debug, Args)]
struct ModelsCommand {
    #[command(subcommand)]
    command: ModelsSubcommand,
}

#[derive(Debug, Subcommand)]
enum ModelsSubcommand {
    List,
    Scan,
    Load(ModelsLoadArgs),
    Unload(ModelsUnloadArgs),
    Use(ModelsUseArgs),
    Watch,
    Download(ModelsDownloadArgs),
    Remove(ModelsRemoveArgs),
}

#[derive(Debug, Args)]
struct ModelsLoadArgs {
    model: String,

    #[arg(long)]
    alias: Option<String>,
}

#[derive(Debug, Args)]
struct ModelsUnloadArgs {
    model: Option<String>,
}

#[derive(Debug, Args)]
struct ModelsUseArgs {
    model: String,
}

#[derive(Debug, Args)]
struct ModelsDownloadArgs {
    repo: String,
}

#[derive(Debug, Args)]
struct ModelsRemoveArgs {
    model: String,
}

#[derive(Debug, Args)]
struct AuthCommand {
    #[command(subcommand)]
    command: AuthSubcommand,
}

#[derive(Debug, Subcommand)]
enum AuthSubcommand {
    Status,
    Enable(AuthEnableArgs),
    Disable,
    SetKey(AuthSetKeyArgs),
}

#[derive(Debug, Args)]
struct AuthEnableArgs {
    #[arg(long, value_enum, default_value_t = AuthMode::Protected)]
    mode: AuthMode,

    #[arg(long)]
    api_key: Option<String>,
}

#[derive(Debug, Args)]
struct AuthSetKeyArgs {
    api_key: String,
}

#[derive(Debug, Args)]
struct ConfigCommand {
    #[command(subcommand)]
    command: ConfigSubcommand,
}

#[derive(Debug, Subcommand)]
enum ConfigSubcommand {
    Get(ConfigGetArgs),
    Set(ConfigSetArgs),
}

#[derive(Debug, Args)]
struct ConfigGetArgs {
    key: String,
}

#[derive(Debug, Args)]
struct ConfigSetArgs {
    key: String,
    value: String,
}

#[derive(Debug, Args)]
struct BenchCommand {
    #[command(subcommand)]
    command: BenchSubcommand,
}

#[derive(Debug, Subcommand)]
enum BenchSubcommand {
    Throughput(BenchThroughputArgs),
    Smoke,
    Suite,
    Compare(BenchCompareArgs),
}

#[derive(Debug, Args)]
struct BenchThroughputArgs {
    #[arg(long, default_value_t = 3)]
    samples: u32,

    #[arg(long, default_value_t = 256)]
    prompt_tokens: u32,

    #[arg(long, default_value_t = 128)]
    completion_tokens: u32,
}

#[derive(Debug, Args)]
struct BenchCompareArgs {
    baseline: String,
    candidate: String,
}

#[derive(Debug, Args)]
struct LogsArgs {
    #[arg(long, default_value_t = 200)]
    lines: u32,
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum CommandPlan {
    SummaryOnly {
        command: String,
    },
    ProbeStatus {
        server: String,
        api_key_present: bool,
        format: OutputFormat,
    },
    ScanModels {
        models_dir: String,
    },
    LoadModel {
        models_dir: String,
        model: String,
        alias: Option<String>,
    },
    ConfigureAuth {
        mode: AuthMode,
        api_key: Option<String>,
    },
    SetConfig {
        key: String,
        value: String,
    },
    RunThroughputBench {
        server: String,
        api_key_present: bool,
        samples: u32,
        prompt_tokens: u32,
        completion_tokens: u32,
    },
    ReadLogs {
        lines: u32,
    },
}

fn main() -> anyhow::Result<()> {
    let cli = Cli::parse();
    println!("CLI scaffold ready for {}", render_command_summary(&cli));
    Ok(())
}

fn render_command_summary(cli: &Cli) -> String {
    format!(
        "command={} server={} format={} models_dir={}",
        command_summary(&cli.command),
        cli.server,
        output_format_name(cli.format),
        cli.models_dir
    )
}

fn plan_command(cli: &Cli) -> CommandPlan {
    match &cli.command {
        Commands::Status => CommandPlan::ProbeStatus {
            server: cli.server.clone(),
            api_key_present: cli.api_key.is_some(),
            format: cli.format,
        },
        Commands::Models(models) => match &models.command {
            ModelsSubcommand::Scan => CommandPlan::ScanModels {
                models_dir: cli.models_dir.clone(),
            },
            ModelsSubcommand::Load(args) => CommandPlan::LoadModel {
                models_dir: cli.models_dir.clone(),
                model: args.model.clone(),
                alias: args.alias.clone(),
            },
            _ => CommandPlan::SummaryOnly {
                command: command_summary(&cli.command),
            },
        },
        Commands::Auth(auth) => match &auth.command {
            AuthSubcommand::Enable(args) => CommandPlan::ConfigureAuth {
                mode: args.mode,
                api_key: args.api_key.clone(),
            },
            _ => CommandPlan::SummaryOnly {
                command: command_summary(&cli.command),
            },
        },
        Commands::Config(config) => match &config.command {
            ConfigSubcommand::Set(args) => CommandPlan::SetConfig {
                key: args.key.clone(),
                value: args.value.clone(),
            },
            _ => CommandPlan::SummaryOnly {
                command: command_summary(&cli.command),
            },
        },
        Commands::Bench(bench) => match &bench.command {
            BenchSubcommand::Throughput(args) => CommandPlan::RunThroughputBench {
                server: cli.server.clone(),
                api_key_present: cli.api_key.is_some(),
                samples: args.samples,
                prompt_tokens: args.prompt_tokens,
                completion_tokens: args.completion_tokens,
            },
            _ => CommandPlan::SummaryOnly {
                command: command_summary(&cli.command),
            },
        },
        Commands::Logs(args) => CommandPlan::ReadLogs { lines: args.lines },
        Commands::Doctor => CommandPlan::SummaryOnly {
            command: command_summary(&cli.command),
        },
    }
}

fn command_summary(command: &Commands) -> String {
    match command {
        Commands::Status => "status".to_owned(),
        Commands::Models(models) => models_command_summary(models),
        Commands::Auth(auth) => auth_command_summary(auth),
        Commands::Config(config) => config_command_summary(config),
        Commands::Bench(bench) => bench_command_summary(bench),
        Commands::Logs(args) => format!("logs lines={}", args.lines),
        Commands::Doctor => "doctor".to_owned(),
    }
}

fn output_format_name(format: OutputFormat) -> &'static str {
    match format {
        OutputFormat::Text => "text",
        OutputFormat::Json => "json",
    }
}

fn models_command_summary(command: &ModelsCommand) -> String {
    match &command.command {
        ModelsSubcommand::List => "models.list".to_owned(),
        ModelsSubcommand::Scan => "models.scan".to_owned(),
        ModelsSubcommand::Load(args) => match &args.alias {
            Some(alias) => format!("models.load model={} alias={}", args.model, alias),
            None => format!("models.load model={}", args.model),
        },
        ModelsSubcommand::Unload(args) => match &args.model {
            Some(model) => format!("models.unload model={}", model),
            None => "models.unload model=active".to_owned(),
        },
        ModelsSubcommand::Use(args) => format!("models.use model={}", args.model),
        ModelsSubcommand::Watch => "models.watch".to_owned(),
        ModelsSubcommand::Download(args) => format!("models.download repo={}", args.repo),
        ModelsSubcommand::Remove(args) => format!("models.remove model={}", args.model),
    }
}

fn auth_command_summary(command: &AuthCommand) -> String {
    match &command.command {
        AuthSubcommand::Status => "auth.status".to_owned(),
        AuthSubcommand::Enable(args) => format!(
            "auth.enable mode={} has_api_key={}",
            auth_mode_name(args.mode),
            args.api_key.is_some()
        ),
        AuthSubcommand::Disable => "auth.disable".to_owned(),
        AuthSubcommand::SetKey(_) => "auth.set-key has_api_key=true".to_owned(),
    }
}

fn auth_mode_name(mode: AuthMode) -> &'static str {
    match mode {
        AuthMode::Off => "off",
        AuthMode::LocalOnly => "local-only",
        AuthMode::Protected => "protected",
    }
}

fn config_command_summary(command: &ConfigCommand) -> String {
    match &command.command {
        ConfigSubcommand::Get(args) => format!("config.get key={}", args.key),
        ConfigSubcommand::Set(args) => format!("config.set key={} value={}", args.key, args.value),
    }
}

fn bench_command_summary(command: &BenchCommand) -> String {
    match &command.command {
        BenchSubcommand::Throughput(args) => format!(
            "bench.throughput samples={} prompt_tokens={} completion_tokens={}",
            args.samples, args.prompt_tokens, args.completion_tokens
        ),
        BenchSubcommand::Smoke => "bench.smoke".to_owned(),
        BenchSubcommand::Suite => "bench.suite".to_owned(),
        BenchSubcommand::Compare(args) => {
            format!(
                "bench.compare baseline={} candidate={}",
                args.baseline, args.candidate
            )
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{plan_command, render_command_summary, AuthMode, Cli, CommandPlan, Commands};
    use clap::Parser;

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_status_command_supports_global_output_flags() {
        let parsed = Cli::try_parse_from(["localllm", "--format", "json", "status"]);
        assert!(parsed.is_ok());
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_models_list_command_is_available() {
        let parsed = Cli::try_parse_from(["localllm", "models", "list"]);
        assert!(parsed.is_ok());
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_models_scan_command_is_available() {
        let parsed = Cli::try_parse_from(["localllm", "models", "scan"]);
        assert!(parsed.is_ok());
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_models_load_command_accepts_model_and_alias() {
        let parsed = Cli::try_parse_from([
            "localllm",
            "models",
            "load",
            "./models/qwen.gguf",
            "--alias",
            "qwen-coder",
        ]);
        assert!(parsed.is_ok());
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_models_watch_command_is_available() {
        let parsed = Cli::try_parse_from(["localllm", "models", "watch"]);
        assert!(parsed.is_ok());
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_models_download_command_accepts_hf_repo() {
        let parsed = Cli::try_parse_from([
            "localllm",
            "models",
            "download",
            "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF",
        ]);
        assert!(parsed.is_ok());
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_auth_enable_command_supports_mode_and_key() {
        let parsed = Cli::try_parse_from([
            "localllm",
            "auth",
            "enable",
            "--mode",
            "protected",
            "--api-key",
            "secret",
        ]);
        assert!(parsed.is_ok());
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_config_set_command_supports_key_value_updates() {
        let parsed =
            Cli::try_parse_from(["localllm", "config", "set", "default_model", "qwen-coder"]);
        assert!(parsed.is_ok());
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_bench_throughput_command_accepts_sample_controls() {
        let parsed = Cli::try_parse_from([
            "localllm",
            "bench",
            "throughput",
            "--samples",
            "5",
            "--prompt-tokens",
            "256",
            "--completion-tokens",
            "128",
        ]);
        assert!(parsed.is_ok());
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_logs_command_is_available() {
        let parsed = Cli::try_parse_from(["localllm", "logs"]);
        assert!(parsed.is_ok());
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_doctor_command_is_available() {
        let parsed = Cli::try_parse_from(["localllm", "doctor"]);
        assert!(parsed.is_ok());
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_status_summary_includes_output_contract() {
        let cli = Cli::try_parse_from(["localllm", "--format", "json", "status"]).unwrap();
        let summary = render_command_summary(&cli);

        assert!(summary.contains("status"));
        assert!(summary.contains("format=json"));
        assert!(summary.contains("models_dir=./models"));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_models_load_summary_includes_model_and_alias() {
        let cli = Cli::try_parse_from([
            "localllm",
            "models",
            "load",
            "./models/qwen.gguf",
            "--alias",
            "qwen-coder",
        ])
        .unwrap();
        let summary = render_command_summary(&cli);

        assert!(summary.contains("models.load"));
        assert!(summary.contains("./models/qwen.gguf"));
        assert!(summary.contains("alias=qwen-coder"));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_auth_enable_summary_includes_mode_and_key_presence() {
        let cli = Cli::try_parse_from([
            "localllm",
            "auth",
            "enable",
            "--mode",
            "local-only",
            "--api-key",
            "secret",
        ])
        .unwrap();
        let summary = render_command_summary(&cli);

        assert!(summary.contains("auth.enable"));
        assert!(summary.contains("mode=local-only"));
        assert!(summary.contains("has_api_key=true"));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_bench_throughput_summary_includes_default_budget() {
        let cli = Cli::try_parse_from(["localllm", "bench", "throughput"]).unwrap();
        let summary = render_command_summary(&cli);

        assert!(summary.contains("bench.throughput"));
        assert!(summary.contains("samples=3"));
        assert!(summary.contains("prompt_tokens=256"));
        assert!(summary.contains("completion_tokens=128"));
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_logs_summary_includes_default_line_limit() {
        let cli = Cli::try_parse_from(["localllm", "logs"]).unwrap();
        let summary = render_command_summary(&cli);

        assert!(summary.contains("logs"));
        assert!(summary.contains("lines=200"));
    }

    #[test]
    fn auth_mode_parser_accepts_local_only() {
        let cli =
            Cli::try_parse_from(["localllm", "auth", "enable", "--mode", "local-only"]).unwrap();

        match cli.command {
            Commands::Auth(auth) => match auth.command {
                super::AuthSubcommand::Enable(args) => assert_eq!(args.mode, AuthMode::LocalOnly),
                other => panic!("unexpected auth command: {:?}", other),
            },
            other => panic!("unexpected command: {:?}", other),
        }
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_status_plan_keeps_server_and_format() {
        let cli = Cli::try_parse_from(["localllm", "--format", "json", "status"]).unwrap();

        assert_eq!(
            plan_command(&cli),
            CommandPlan::ProbeStatus {
                server: "http://127.0.0.1:8080".to_owned(),
                api_key_present: false,
                format: super::OutputFormat::Json,
            }
        );
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_models_scan_plan_uses_models_dir() {
        let cli = Cli::try_parse_from(["localllm", "--models-dir", "D:/models", "models", "scan"])
            .unwrap();

        assert_eq!(
            plan_command(&cli),
            CommandPlan::ScanModels {
                models_dir: "D:/models".to_owned(),
            }
        );
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_models_load_plan_preserves_alias_and_models_dir() {
        let cli = Cli::try_parse_from([
            "localllm",
            "--models-dir",
            "D:/models",
            "models",
            "load",
            "qwen.gguf",
            "--alias",
            "qwen-coder",
        ])
        .unwrap();

        assert_eq!(
            plan_command(&cli),
            CommandPlan::LoadModel {
                models_dir: "D:/models".to_owned(),
                model: "qwen.gguf".to_owned(),
                alias: Some("qwen-coder".to_owned()),
            }
        );
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_auth_enable_plan_preserves_mode_and_key() {
        let cli = Cli::try_parse_from([
            "localllm",
            "auth",
            "enable",
            "--mode",
            "protected",
            "--api-key",
            "secret",
        ])
        .unwrap();

        assert_eq!(
            plan_command(&cli),
            CommandPlan::ConfigureAuth {
                mode: AuthMode::Protected,
                api_key: Some("secret".to_owned()),
            }
        );
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_config_set_plan_preserves_key_value() {
        let cli = Cli::try_parse_from(["localllm", "config", "set", "default_model", "qwen-coder"])
            .unwrap();

        assert_eq!(
            plan_command(&cli),
            CommandPlan::SetConfig {
                key: "default_model".to_owned(),
                value: "qwen-coder".to_owned(),
            }
        );
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_bench_throughput_plan_preserves_workload() {
        let cli = Cli::try_parse_from([
            "localllm",
            "--api-key",
            "secret",
            "bench",
            "throughput",
            "--samples",
            "5",
            "--prompt-tokens",
            "512",
            "--completion-tokens",
            "64",
        ])
        .unwrap();

        assert_eq!(
            plan_command(&cli),
            CommandPlan::RunThroughputBench {
                server: "http://127.0.0.1:8080".to_owned(),
                api_key_present: true,
                samples: 5,
                prompt_tokens: 512,
                completion_tokens: 64,
            }
        );
    }

    #[test]
    #[ignore = "autoresearch benchmark"]
    fn cli_surface_logs_plan_preserves_line_limit() {
        let cli = Cli::try_parse_from(["localllm", "logs", "--lines", "50"]).unwrap();

        assert_eq!(plan_command(&cli), CommandPlan::ReadLogs { lines: 50 });
    }
}