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
140
141
//! 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, env = "LLMSIM_HOST")]
host: Option<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: Option<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;
if let Some(host) = host {
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(())
}