use std::net::{SocketAddr, ToSocketAddrs};
use std::path::PathBuf;
use std::process::ExitCode;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use std::fs::File;
use std::io::BufReader;
use llmp::proxy::{self, BodyMode, JsonlSink, ProxyConfig, Sink};
use llmp::tokenize;
use llmp::{Config, Record, read_jsonl};
const DEFAULT_LISTEN: &str = "127.0.0.1:8080";
const DEFAULT_MAX_BODY: usize = 64 * 1024 * 1024;
const USAGE: &str = "\
usage:
llmp proxy --upstream HOST:PORT [--listen HOST:PORT] [--out FILE] [--digest] [--max-body BYTES]
llmp analyze FILE... [--json] [--tokenize URL] [--gap-seconds N]
llmp --version
";
fn main() -> ExitCode {
let args: Vec<String> = std::env::args().skip(1).collect();
match args.first().map(String::as_str) {
Some("--version" | "-V") => {
println!("llmp {}", env!("CARGO_PKG_VERSION"));
ExitCode::SUCCESS
}
Some("--help" | "-h") => {
print!("{USAGE}");
ExitCode::SUCCESS
}
Some("proxy") => proxy(&args[1..]),
Some("analyze") => analyze(&args[1..]),
_ => {
eprint!("{USAGE}");
ExitCode::from(2)
}
}
}
fn proxy(args: &[String]) -> ExitCode {
let (config, out) = match proxy_options(args) {
Ok(options) => options,
Err(message) => return usage_error(&message),
};
let sink = match JsonlSink::create(&out) {
Ok(sink) => Arc::new(sink),
Err(e) => {
eprintln!("error: cannot open {}: {e}", out.display());
return ExitCode::from(1);
}
};
let seen = Arc::new(AtomicU64::new(0));
let forward = {
let sink = Arc::clone(&sink);
let seen = Arc::clone(&seen);
move |record: Record| {
seen.fetch_add(1, Ordering::Relaxed);
sink.record(record);
}
};
println!(
"llmp: {} -> {}, recording to {}",
config.listen,
config.upstream,
out.display()
);
let served = proxy::serve_blocking(config, forward);
let written = match Arc::try_unwrap(sink) {
Ok(sink) => sink.close(),
Err(_) => Ok(seen.load(Ordering::Relaxed)),
};
if let Err(e) = served {
eprintln!("error: {e}");
return ExitCode::from(1);
}
match written {
Ok(count) => {
println!("llmp: {count} records written to {}", out.display());
ExitCode::SUCCESS
}
Err(e) => {
eprintln!("error: writing {}: {e}", out.display());
ExitCode::from(1)
}
}
}
fn proxy_options(args: &[String]) -> Result<(ProxyConfig, PathBuf), String> {
let mut upstream = None;
let mut listen = None;
let mut out = None;
let mut body_mode = BodyMode::Full;
let mut max_body_bytes = DEFAULT_MAX_BODY;
let mut rest = args.iter();
while let Some(arg) = rest.next() {
let mut value = || {
rest.next()
.map(String::as_str)
.ok_or_else(|| format!("{arg} needs a value"))
};
match arg.as_str() {
"--upstream" => upstream = Some(resolve(value()?)?),
"--listen" => listen = Some(resolve(value()?)?),
"--out" => out = Some(PathBuf::from(value()?)),
"--max-body" => {
let raw = value()?;
max_body_bytes = raw
.parse()
.map_err(|_| format!("--max-body: not a byte count: {raw}"))?;
}
"--digest" => body_mode = BodyMode::Digest,
other => return Err(format!("unknown argument: {other}")),
}
}
let upstream = upstream.ok_or("--upstream is required")?;
let listen = match listen {
Some(listen) => listen,
None => resolve(DEFAULT_LISTEN)?,
};
let out = out.unwrap_or_else(|| {
let seconds = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |d| d.as_secs());
PathBuf::from(format!("llmp-{seconds}.jsonl"))
});
let config = ProxyConfig {
listen,
upstream,
body_mode,
max_body_bytes,
};
Ok((config, out))
}
fn resolve(address: &str) -> Result<SocketAddr, String> {
let addresses: Vec<SocketAddr> = address
.to_socket_addrs()
.map_err(|e| format!("{address}: {e}"))?
.collect();
addresses
.iter()
.find(|a| a.is_ipv4())
.or_else(|| addresses.first())
.copied()
.ok_or_else(|| format!("{address}: no address"))
}
fn usage_error(message: &str) -> ExitCode {
eprintln!("error: {message}");
eprint!("{USAGE}");
ExitCode::from(2)
}
struct AnalyzeOptions {
files: Vec<PathBuf>,
json: bool,
tokenize: Option<String>,
config: Config,
}
fn analyze(args: &[String]) -> ExitCode {
let options = match analyze_options(args) {
Ok(options) => options,
Err(message) => return usage_error(&message),
};
let mut records = Vec::new();
let mut bad_lines = 0u64;
for (index, path) in options.files.iter().enumerate() {
let file = match File::open(path) {
Ok(file) => file,
Err(e) => {
eprintln!("error: cannot open {}: {e}", path.display());
return ExitCode::from(1);
}
};
let read = read_jsonl(BufReader::new(file));
for bad in &read.bad_lines {
eprintln!("warning: {}:{}: {}", path.display(), bad.line, bad.error);
}
bad_lines += read.bad_lines.len() as u64;
records.extend(read.records.into_iter().map(|record| (index, record)));
}
records.sort_by_key(|(index, record)| (record.t_start_ms, *index, record.id));
let renumber = options.files.len() > 1;
let mut records: Vec<Record> = records
.into_iter()
.enumerate()
.map(|(position, (_, mut record))| {
if renumber {
record.id = position as u64 + 1;
}
record
})
.collect();
if let Some(server) = &options.tokenize {
let result = tokenize::enrich(&mut records, server);
for warning in &result.warnings {
eprintln!("warning: tokenize: {warning}");
}
eprintln!(
"llmp: counted tokens for {} texts at {server}",
result.counted
);
}
let mut report = llmp::analyze(&records, &options.config);
report.coverage.bad_lines = bad_lines;
if options.json {
match serde_json::to_string_pretty(&report) {
Ok(json) => println!("{json}"),
Err(e) => {
eprintln!("error: {e}");
return ExitCode::from(1);
}
}
} else {
print!("{}", report.summary());
}
ExitCode::SUCCESS
}
fn analyze_options(args: &[String]) -> Result<AnalyzeOptions, String> {
let mut options = AnalyzeOptions {
files: Vec::new(),
json: false,
tokenize: None,
config: Config::default(),
};
let mut rest = args.iter();
while let Some(arg) = rest.next() {
let mut value = || {
rest.next()
.map(String::as_str)
.ok_or_else(|| format!("{arg} needs a value"))
};
match arg.as_str() {
"--json" => options.json = true,
"--tokenize" => options.tokenize = Some(value()?.to_string()),
"--gap-seconds" => {
let raw = value()?;
let seconds: u64 = raw
.parse()
.map_err(|_| format!("--gap-seconds: not a number of seconds: {raw}"))?;
options.config.session_gap_ms = seconds.saturating_mul(1000);
}
flag if flag.starts_with("--") => return Err(format!("unknown argument: {flag}")),
file => options.files.push(PathBuf::from(file)),
}
}
if options.files.is_empty() {
return Err("analyze needs at least one recording".into());
}
Ok(options)
}