llmsim 0.3.0

LLM Traffic Simulator - A lightweight, high-performance LLM API simulator
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

//! LLMSim CLI - LLM Traffic Simulator
//!
//! Usage:
//!   llmsim serve [OPTIONS]    Start the HTTP server
//!
//! Examples:
//!   llmsim serve --port 8080
//!   llmsim serve --config config.toml
//!   llmsim serve --generator echo --target-tokens 50
//!   llmsim serve --tui              # Start with real-time stats dashboard

use clap::{Parser, Subcommand};
use llmsim::cli::{Config, ConfigError};
#[cfg(feature = "tui")]
use llmsim::tui::{run_dashboard, DashboardConfig};

#[derive(Parser)]
#[command(name = "llmsim")]
#[command(author, version, about = "LLM Traffic Simulator", long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Start the LLMSim HTTP server
    Serve {
        /// Configuration file path (TOML)
        #[arg(short, long)]
        config: Option<String>,

        /// Port to listen on
        #[arg(short, long, default_value = "8080")]
        port: u16,

        /// Host to bind to
        #[arg(long, default_value = "0.0.0.0")]
        host: String,

        /// Response generator (lorem, echo, random, fixed:text)
        #[arg(long, default_value = "lorem")]
        generator: String,

        /// Target number of tokens in responses
        #[arg(long, default_value = "100")]
        target_tokens: usize,

        /// Show real-time stats dashboard (TUI)
        ///
        /// Requires building with `--features tui`.
        #[arg(long)]
        tui: bool,
    },
}

fn build_config(
    config_file: Option<String>,
    port: u16,
    host: String,
    generator: String,
    target_tokens: usize,
) -> Result<Config, ConfigError> {
    let mut config = if let Some(path) = config_file {
        Config::from_file(&path)?
    } else {
        Config::default()
    };

    // Override with CLI arguments
    config.server.port = port;
    config.server.host = host;
    config.response.generator = generator;
    config.response.target_tokens = target_tokens;

    Ok(config)
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let cli = Cli::parse();

    match cli.command {
        Commands::Serve {
            config,
            port,
            host,
            generator,
            target_tokens,
            tui,
        } => {
            let config = build_config(config, port, host.clone(), generator, target_tokens)?;

            if tui {
                #[cfg(not(feature = "tui"))]
                {
                    return Err(
                        "the --tui flag requires building llmsim with --features tui".into(),
                    );
                }

                #[cfg(feature = "tui")]
                {
                    // Run server and TUI concurrently
                    let stats = llmsim::new_shared_stats();
                    let server_url = format!("http://127.0.0.1:{}", port);

                    let dashboard_config = DashboardConfig {
                        server_url,
                        refresh_ms: 200,
                    };

                    // Run both concurrently - TUI exit will shut down the app
                    tokio::select! {
                        result = llmsim::cli::run_server_with_stats(config, stats) => {
                            result?;
                        }
                        result = run_dashboard(dashboard_config) => {
                            result?;
                        }
                    }
                }
            } else {
                // Initialize tracing for server-only mode
                tracing_subscriber::fmt()
                    .with_env_filter(
                        tracing_subscriber::EnvFilter::from_default_env()
                            .add_directive("llmsim=info".parse().unwrap())
                            .add_directive("tower_http=debug".parse().unwrap()),
                    )
                    .init();

                llmsim::cli::run_server(config).await?;
            }
        }
    }

    Ok(())
}