systemprompt-cli 0.1.22

systemprompt.io OS - CLI for agent orchestration, AI operations, and system management
Documentation
use anyhow::{Result, bail};
use systemprompt_logging::CliService;
use systemprompt_models::ProfileBootstrap;

use super::SetArgs;
use super::helpers::{
    get_endpoint_rate, get_tier_multiplier, load_profile_for_edit, save_profile, set_endpoint_rate,
    set_tier_multiplier,
};
use crate::CliConfig;
use crate::cli_settings::OutputFormat;
use crate::shared::{CommandResult, render_result};

use super::super::types::{RateLimitStatusOutput, SetRateLimitOutput};

pub fn execute_set(args: &SetArgs, config: &CliConfig) -> Result<()> {
    if args.endpoint.is_some() && args.rate.is_none() {
        bail!("--rate is required when --endpoint is specified");
    }
    if args.rate.is_some() && args.endpoint.is_none() {
        bail!("--endpoint is required when --rate is specified");
    }
    if args.tier.is_some() && args.multiplier.is_none() {
        bail!("--multiplier is required when --tier is specified");
    }
    if args.multiplier.is_some() && args.tier.is_none() {
        bail!("--tier is required when --multiplier is specified");
    }
    if args.endpoint.is_none() && args.tier.is_none() && args.burst.is_none() {
        bail!("Must specify one of: --endpoint with --rate, --tier with --multiplier, or --burst");
    }

    let profile_path = ProfileBootstrap::get_path()?;
    let mut profile = load_profile_for_edit(profile_path)?;
    let limits = &mut profile.rate_limits;

    let output = if let (Some(endpoint), Some(rate)) = (&args.endpoint, args.rate) {
        let old_value = get_endpoint_rate(limits, endpoint)?;
        set_endpoint_rate(limits, endpoint, rate)?;
        SetRateLimitOutput {
            field: format!("{}_per_second", endpoint),
            old_value: old_value.to_string(),
            new_value: rate.to_string(),
            message: format!("Updated {} rate: {} -> {}/s", endpoint, old_value, rate),
        }
    } else if let (Some(tier), Some(multiplier)) = (&args.tier, args.multiplier) {
        let old_value = get_tier_multiplier(&limits.tier_multipliers, tier)?;
        set_tier_multiplier(&mut limits.tier_multipliers, tier, multiplier)?;
        SetRateLimitOutput {
            field: format!("tier_multipliers.{}", tier),
            old_value: format!("{:.1}", old_value),
            new_value: format!("{:.1}", multiplier),
            message: format!(
                "Updated {} tier multiplier: {:.1}x -> {:.1}x",
                tier, old_value, multiplier
            ),
        }
    } else if let Some(burst) = args.burst {
        let old_value = limits.burst_multiplier;
        limits.burst_multiplier = burst;
        SetRateLimitOutput {
            field: "burst_multiplier".to_string(),
            old_value: old_value.to_string(),
            new_value: burst.to_string(),
            message: format!("Updated burst multiplier: {}x -> {}x", old_value, burst),
        }
    } else {
        bail!("Invalid arguments");
    };

    save_profile(&profile, profile_path)?;
    render_result(&CommandResult::text(output).with_title("Rate Limit Updated"));

    if config.output_format() == OutputFormat::Table {
        CliService::warning("Restart services for changes to take effect");
    }

    Ok(())
}

pub fn execute_enable(config: &CliConfig) -> Result<()> {
    let profile_path = ProfileBootstrap::get_path()?;
    let mut profile = load_profile_for_edit(profile_path)?;

    if !profile.rate_limits.disabled {
        let output = RateLimitStatusOutput {
            enabled: true,
            message: "Rate limiting is already enabled".to_string(),
        };
        render_result(&CommandResult::text(output).with_title("Rate Limit Status"));
        return Ok(());
    }

    profile.rate_limits.disabled = false;
    save_profile(&profile, profile_path)?;

    let output = RateLimitStatusOutput {
        enabled: true,
        message: "Rate limiting enabled".to_string(),
    };
    render_result(&CommandResult::text(output).with_title("Rate Limit Status"));

    if config.output_format() == OutputFormat::Table {
        CliService::warning("Restart services for changes to take effect");
    }

    Ok(())
}

pub fn execute_disable(config: &CliConfig) -> Result<()> {
    let profile_path = ProfileBootstrap::get_path()?;
    let mut profile = load_profile_for_edit(profile_path)?;

    if profile.rate_limits.disabled {
        let output = RateLimitStatusOutput {
            enabled: false,
            message: "Rate limiting is already disabled".to_string(),
        };
        render_result(&CommandResult::text(output).with_title("Rate Limit Status"));
        return Ok(());
    }

    profile.rate_limits.disabled = true;
    save_profile(&profile, profile_path)?;

    let output = RateLimitStatusOutput {
        enabled: false,
        message: "Rate limiting disabled".to_string(),
    };
    render_result(&CommandResult::text(output).with_title("Rate Limit Status"));

    if config.output_format() == OutputFormat::Table {
        CliService::warning("Restart services for changes to take effect");
    }

    Ok(())
}