roy_cli/
lib.rs

1// Copyright 2025 Massimiliano Pippi
2// SPDX-License-Identifier: MIT
3
4use axum::{routing::post, Router};
5use clap::Parser;
6use clap_verbosity_flag::Verbosity;
7use colored::Colorize;
8use std::net::{IpAddr, SocketAddr};
9
10pub mod chat_completion;
11pub mod server_state;
12use server_state::ServerState;
13
14#[derive(Parser, Clone)]
15#[command(name = "roy")]
16#[command(version = env!("CARGO_PKG_VERSION"))]
17#[command(
18    about = "A HTTP server compatible with the OpenAI platform format that simulates errors and rate limit data"
19)]
20pub struct Args {
21    #[command(flatten)]
22    pub verbosity: Verbosity,
23
24    #[arg(long, help = "Port to listen on", default_value = "8000")]
25    pub port: u16,
26
27    #[arg(long, help = "Address to listen on", default_value = "127.0.0.1")]
28    pub address: IpAddr,
29
30    #[arg(
31        long,
32        help = "Length of response (fixed number or range like '10:100')",
33        default_value = "250"
34    )]
35    pub response_length: Option<String>,
36
37    #[arg(long, help = "HTTP error code to return")]
38    pub error_code: Option<u16>,
39
40    #[arg(long, help = "Error rate percentage (0-100)")]
41    pub error_rate: Option<u32>,
42
43    #[arg(
44        long,
45        help = "Maximum number of requests per minute",
46        default_value = "500"
47    )]
48    pub rpm: u32,
49
50    #[arg(
51        long,
52        help = "Maximum number of tokens per minute",
53        default_value = "30000"
54    )]
55    pub tpm: u32,
56}
57
58pub async fn run(args: Args) -> anyhow::Result<()> {
59    let state = ServerState::new(args.clone());
60
61    let app = Router::new()
62        .route(
63            "/v1/chat/completions",
64            post(chat_completion::chat_completions),
65        )
66        .with_state(state);
67
68    let addr = SocketAddr::new(args.address, args.port);
69    let listener = tokio::net::TcpListener::bind(addr).await?;
70
71    println!(
72        "Roy server running on {}",
73        format!("http://{}", addr).blue()
74    );
75    axum::serve(listener, app).await?;
76
77    Ok(())
78}
roy_cli/lib.rs

roy_cli/
lib.rs
