feat: implement clap CLI subcommands and terminal UX
This commit is contained in:
parent
e67a48b81b
commit
fe3fc74e1c
2 changed files with 606 additions and 4 deletions
89
src/cli.rs
Normal file
89
src/cli.rs
Normal file
|
|
@ -0,0 +1,89 @@
|
||||||
|
use clap::{Args, Parser, Subcommand};
|
||||||
|
use std::path::PathBuf;
|
||||||
|
|
||||||
|
#[derive(Parser, Debug)]
|
||||||
|
#[command(
|
||||||
|
name = "brokecrew",
|
||||||
|
version,
|
||||||
|
about = "A Rust-native multi-model AI team CLI with resilient OpenRouter fallback routing",
|
||||||
|
long_about = "BrokeCrew lets you interact with a coordinated team of AI personas backed by OpenRouter models.\nFeatures two-stage routing, automatic candidate fallbacks, budget tracking, and multi-agent handoff pipelines."
|
||||||
|
)]
|
||||||
|
pub struct Cli {
|
||||||
|
#[arg(long, global = true, help = "Path to config.toml file")]
|
||||||
|
pub config_file: Option<PathBuf>,
|
||||||
|
|
||||||
|
#[arg(long, global = true, help = "Path to personas.toml file")]
|
||||||
|
pub personas_file: Option<PathBuf>,
|
||||||
|
|
||||||
|
#[arg(long, global = true, env = "OPENROUTER_API_KEY", help = "OpenRouter API Key (overrides keyring & env)")]
|
||||||
|
pub api_key: Option<String>,
|
||||||
|
|
||||||
|
#[arg(long, global = true, help = "Output structured logs in JSON format to stderr")]
|
||||||
|
pub json_logs: bool,
|
||||||
|
|
||||||
|
#[arg(long, short, global = true, help = "Enable verbose debug logs")]
|
||||||
|
pub verbose: bool,
|
||||||
|
|
||||||
|
#[command(subcommand)]
|
||||||
|
pub command: Option<Commands>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Subcommand, Debug)]
|
||||||
|
pub enum Commands {
|
||||||
|
#[command(about = "Chat with a persona or multi-agent pipeline")]
|
||||||
|
Chat(ChatArgs),
|
||||||
|
|
||||||
|
#[command(about = "Manage configuration and credentials")]
|
||||||
|
Config(ConfigArgs),
|
||||||
|
|
||||||
|
#[command(about = "List and inspect personas and pipelines")]
|
||||||
|
Personas(PersonasArgs),
|
||||||
|
|
||||||
|
#[command(about = "Display live model registry and rate-limit budget status")]
|
||||||
|
Status,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Args, Debug, Default)]
|
||||||
|
pub struct ChatArgs {
|
||||||
|
#[arg(help = "The message or prompt to send. If omitted, starts interactive chat REPL.")]
|
||||||
|
pub prompt: Option<String>,
|
||||||
|
|
||||||
|
#[arg(short, long, help = "Directly specify target persona (skips router)")]
|
||||||
|
pub persona: Option<String>,
|
||||||
|
|
||||||
|
#[arg(long, help = "Run a multi-agent pipeline (e.g. build-and-review)")]
|
||||||
|
pub pipeline: Option<String>,
|
||||||
|
|
||||||
|
#[arg(long, help = "Allow experimental stealth model candidates")]
|
||||||
|
pub allow_stealth: bool,
|
||||||
|
|
||||||
|
#[arg(long, help = "Force escalation to paid model candidates")]
|
||||||
|
pub escalate: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Args, Debug)]
|
||||||
|
pub struct ConfigArgs {
|
||||||
|
#[command(subcommand)]
|
||||||
|
pub action: ConfigSubcommand,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Subcommand, Debug)]
|
||||||
|
pub enum ConfigSubcommand {
|
||||||
|
#[command(about = "Securely set and save OpenRouter API key to OS credential store")]
|
||||||
|
SetKey {
|
||||||
|
#[arg(help = "Optional key directly passed; if omitted, securely prompted")]
|
||||||
|
key: Option<String>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Args, Debug)]
|
||||||
|
pub struct PersonasArgs {
|
||||||
|
#[command(subcommand)]
|
||||||
|
pub action: PersonasSubcommand,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Subcommand, Debug)]
|
||||||
|
pub enum PersonasSubcommand {
|
||||||
|
#[command(about = "List all configured personas, candidates, triggers, and pipelines")]
|
||||||
|
List,
|
||||||
|
}
|
||||||
521
src/main.rs
521
src/main.rs
|
|
@ -1,9 +1,522 @@
|
||||||
//! BrokeCrew CLI Entrypoint
|
mod cli;
|
||||||
//!
|
|
||||||
//! A thin executable layer over the `brokecrew` library.
|
use brokecrew::budget::BudgetTracker;
|
||||||
|
use brokecrew::client::OpenRouterClient;
|
||||||
|
use brokecrew::config::AppConfig;
|
||||||
|
use brokecrew::persona::Tier;
|
||||||
|
use brokecrew::pipeline::{PipelineConfig, PipelineEvent, PipelineExecutor};
|
||||||
|
use brokecrew::registry::ModelRegistry;
|
||||||
|
use brokecrew::router::{RouteStage, Router};
|
||||||
|
use clap::Parser;
|
||||||
|
use cli::{ChatArgs, Cli, Commands, ConfigSubcommand, PersonasSubcommand};
|
||||||
|
use indicatif::{ProgressBar, ProgressStyle};
|
||||||
|
use owo_colors::OwoColorize;
|
||||||
|
use std::io::{self, BufRead, Write};
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::time::Duration;
|
||||||
|
use tokio::sync::mpsc;
|
||||||
|
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
|
||||||
|
|
||||||
|
fn init_logging(json_logs: bool, verbose: bool) {
|
||||||
|
let filter = if verbose {
|
||||||
|
EnvFilter::new("brokecrew=debug,reqwest=debug")
|
||||||
|
} else {
|
||||||
|
EnvFilter::new("brokecrew=info")
|
||||||
|
};
|
||||||
|
|
||||||
|
if json_logs {
|
||||||
|
tracing_subscriber::registry()
|
||||||
|
.with(filter)
|
||||||
|
.with(tracing_subscriber::fmt::layer().json())
|
||||||
|
.init();
|
||||||
|
} else {
|
||||||
|
tracing_subscriber::registry()
|
||||||
|
.with(filter)
|
||||||
|
.with(
|
||||||
|
tracing_subscriber::fmt::layer()
|
||||||
|
.with_writer(io::stderr)
|
||||||
|
.without_time()
|
||||||
|
.with_target(false),
|
||||||
|
)
|
||||||
|
.init();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_persona_color_tag(persona: &str) -> String {
|
||||||
|
match persona.to_lowercase().as_str() {
|
||||||
|
"coder" => format!("[{}]", persona).bright_green().bold().to_string(),
|
||||||
|
"reviewer" => format!("[{}]", persona).bright_yellow().bold().to_string(),
|
||||||
|
"architect" => format!("[{}]", persona).bright_blue().bold().to_string(),
|
||||||
|
"general" | "assistant" => format!("[{}]", persona).bright_magenta().bold().to_string(),
|
||||||
|
_ => format!("[{}]", persona).bright_cyan().bold().to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn create_spinner(message: &'static str) -> ProgressBar {
|
||||||
|
let pb = ProgressBar::new_spinner();
|
||||||
|
pb.set_style(
|
||||||
|
ProgressStyle::default_spinner()
|
||||||
|
.tick_chars("⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏ ")
|
||||||
|
.template("{spinner:.cyan} {msg}")
|
||||||
|
.expect("valid template"),
|
||||||
|
);
|
||||||
|
pb.set_message(message);
|
||||||
|
pb.enable_steady_tick(Duration::from_millis(80));
|
||||||
|
pb
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::main]
|
#[tokio::main]
|
||||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
println!("BrokeCrew CLI v{}", env!("CARGO_PKG_VERSION"));
|
let cli = Cli::parse();
|
||||||
|
init_logging(cli.json_logs, cli.verbose);
|
||||||
|
|
||||||
|
let config = AppConfig::load_from_paths(
|
||||||
|
cli.config_file.as_deref(),
|
||||||
|
cli.personas_file.as_deref(),
|
||||||
|
cli.api_key.clone(),
|
||||||
|
)
|
||||||
|
.map_err(|e| format!("Configuration error: {e}"))?;
|
||||||
|
|
||||||
|
let budget_path = config.config_dir.join("budget.json");
|
||||||
|
let mut budget_tracker = BudgetTracker::new(Some(budget_path), config.credit_tier);
|
||||||
|
let model_registry = Arc::new(ModelRegistry::new());
|
||||||
|
|
||||||
|
// Check subcommand
|
||||||
|
match cli.command {
|
||||||
|
Some(Commands::Config(config_args)) => match config_args.action {
|
||||||
|
ConfigSubcommand::SetKey { key } => {
|
||||||
|
let api_key = match key {
|
||||||
|
Some(k) if !k.trim().is_empty() => k.trim().to_string(),
|
||||||
|
_ => {
|
||||||
|
eprint!("Enter OpenRouter API Key: ");
|
||||||
|
io::stderr().flush()?;
|
||||||
|
let mut input = String::new();
|
||||||
|
io::stdin().read_line(&mut input)?;
|
||||||
|
let trimmed = input.trim().to_string();
|
||||||
|
if trimmed.is_empty() {
|
||||||
|
eprintln!("{}", "Error: API key cannot be empty.".red().bold());
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
trimmed
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
AppConfig::save_api_key_to_keyring(&api_key)
|
||||||
|
.map_err(|e| format!("Failed to save API key to OS credential store: {e}"))?;
|
||||||
|
println!("{}", "✔ OpenRouter API key saved securely in OS credential store.".bright_green().bold());
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
},
|
||||||
|
|
||||||
|
Some(Commands::Personas(personas_args)) => match personas_args.action {
|
||||||
|
PersonasSubcommand::List => {
|
||||||
|
println!("\n{}", "BrokeCrew Personas:".bold().underline());
|
||||||
|
for persona in config.persona_registry.list() {
|
||||||
|
let badge = get_persona_color_tag(&persona.name);
|
||||||
|
println!("\n {} - {}", badge, persona.system_prompt.dimmed());
|
||||||
|
if !persona.trigger_words.is_empty() {
|
||||||
|
println!(" Triggers: {}", persona.trigger_words.join(", ").italic());
|
||||||
|
}
|
||||||
|
println!(" Candidates:");
|
||||||
|
for (i, c) in persona.candidates.iter().enumerate() {
|
||||||
|
let tier_tag = match c.tier {
|
||||||
|
Tier::Free => "FREE".bright_green().to_string(),
|
||||||
|
Tier::Paid => "PAID".bright_yellow().to_string(),
|
||||||
|
};
|
||||||
|
let stealth_tag = if c.is_stealth {
|
||||||
|
format!(" {}", "[stealth]".bright_magenta())
|
||||||
|
} else {
|
||||||
|
String::new()
|
||||||
|
};
|
||||||
|
println!(" {}. {} ({}){}", i + 1, c.id.bold(), tier_tag, stealth_tag);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if !config.pipelines.is_empty() {
|
||||||
|
println!("\n{}", "Configured Pipelines:".bold().underline());
|
||||||
|
for pipe in &config.pipelines {
|
||||||
|
let steps_str: Vec<String> = pipe
|
||||||
|
.steps
|
||||||
|
.iter()
|
||||||
|
.map(|s| {
|
||||||
|
if s.pass_previous_output {
|
||||||
|
format!("{} (with context)", s.persona.bold())
|
||||||
|
} else {
|
||||||
|
s.persona.bold().to_string()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
println!(" • {} ➔ {}", pipe.name.bright_cyan().bold(), steps_str.join(" ➔ "));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
println!();
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
},
|
||||||
|
|
||||||
|
Some(Commands::Status) => {
|
||||||
|
println!("\n{}", "BrokeCrew Status & Rate Limits:".bold().underline());
|
||||||
|
// Fetch live registry for status count
|
||||||
|
let spinner = create_spinner("Fetching live OpenRouter model catalog...");
|
||||||
|
let model_count = match model_registry.fetch_live(&config.openrouter_base_url).await {
|
||||||
|
Ok(count) => {
|
||||||
|
spinner.finish_and_clear();
|
||||||
|
count
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
spinner.finish_and_clear();
|
||||||
|
println!(" Live catalog: {}", format!("Offline / Error: {e}").red());
|
||||||
|
0
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
if model_count > 0 {
|
||||||
|
println!(" Live catalog: {} available models cached", model_count.to_string().bright_green().bold());
|
||||||
|
}
|
||||||
|
|
||||||
|
match budget_tracker.check_budget() {
|
||||||
|
Ok(status) => {
|
||||||
|
let pct = (status.daily_requests * 100) / status.daily_limit;
|
||||||
|
let daily_color = if pct >= 80 {
|
||||||
|
format!("{}/{} ({}%)", status.daily_requests, status.daily_limit, pct).yellow().bold().to_string()
|
||||||
|
} else {
|
||||||
|
format!("{}/{} ({}%)", status.daily_requests, status.daily_limit, pct).green().bold().to_string()
|
||||||
|
};
|
||||||
|
|
||||||
|
println!(" Account tier: {:?}", config.credit_tier);
|
||||||
|
println!(" Rolling 1-min requests: {}/{}", status.minute_requests.to_string().bold(), status.minute_limit);
|
||||||
|
println!(" Daily requests today: {}", daily_color);
|
||||||
|
if status.approaching_daily_limit {
|
||||||
|
println!(" {}", "⚠ Notice: Approaching daily free-tier limit. Non-urgent calls may be deferred.".yellow().bold());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
println!(" Budget status: {}", e.to_string().red().bold());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
println!();
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
Some(Commands::Chat(args)) => {
|
||||||
|
run_chat(args, &config, &mut budget_tracker, model_registry).await?;
|
||||||
|
}
|
||||||
|
|
||||||
|
None => {
|
||||||
|
// Default to chat if no subcommand passed
|
||||||
|
run_chat(ChatArgs::default(), &config, &mut budget_tracker, model_registry).await?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run_chat(
|
||||||
|
args: ChatArgs,
|
||||||
|
config: &AppConfig,
|
||||||
|
budget_tracker: &mut BudgetTracker,
|
||||||
|
model_registry: Arc<ModelRegistry>,
|
||||||
|
) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
|
let api_key = match config.api_key.clone() {
|
||||||
|
Some(k) if !k.trim().is_empty() => k,
|
||||||
|
_ => {
|
||||||
|
eprintln!(
|
||||||
|
"{}\nRun: {}\nOr set: {}",
|
||||||
|
"Error: OpenRouter API key not configured.".red().bold(),
|
||||||
|
"brokecrew config set-key".bright_green().bold(),
|
||||||
|
"export OPENROUTER_API_KEY=\"sk-or-v1-...\"".cyan()
|
||||||
|
);
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let client = OpenRouterClient::with_base_url(api_key, config.openrouter_base_url.clone());
|
||||||
|
|
||||||
|
// Start background sync for model registry
|
||||||
|
let reg_clone = model_registry.clone();
|
||||||
|
let base_url = config.openrouter_base_url.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let _ = reg_clone.fetch_live(&base_url).await;
|
||||||
|
});
|
||||||
|
|
||||||
|
let router = Router::new(
|
||||||
|
config.default_persona.clone(),
|
||||||
|
config.router_model.clone(),
|
||||||
|
config.complexity_threshold,
|
||||||
|
);
|
||||||
|
|
||||||
|
// If prompt passed directly as CLI argument, execute single-shot
|
||||||
|
if let Some(ref prompt) = args.prompt {
|
||||||
|
execute_prompt(
|
||||||
|
prompt,
|
||||||
|
&args,
|
||||||
|
config,
|
||||||
|
budget_tracker,
|
||||||
|
&model_registry,
|
||||||
|
&router,
|
||||||
|
&client,
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Otherwise, run interactive REPL loop
|
||||||
|
println!("\n{}", "BrokeCrew Interactive Team Chat".bright_cyan().bold());
|
||||||
|
println!("{}", "Type your request and press Enter. Type 'exit' or Ctrl+C to quit.\n".dimmed());
|
||||||
|
|
||||||
|
let stdin = io::stdin();
|
||||||
|
let mut reader = stdin.lock();
|
||||||
|
|
||||||
|
loop {
|
||||||
|
print!("{} ", "you ❯".bright_white().bold());
|
||||||
|
io::stdout().flush()?;
|
||||||
|
|
||||||
|
let mut line = String::new();
|
||||||
|
if reader.read_line(&mut line)? == 0 {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
let trimmed = line.trim();
|
||||||
|
if trimmed.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if trimmed.eq_ignore_ascii_case("exit") || trimmed.eq_ignore_ascii_case("quit") {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
execute_prompt(
|
||||||
|
trimmed,
|
||||||
|
&args,
|
||||||
|
config,
|
||||||
|
budget_tracker,
|
||||||
|
&model_registry,
|
||||||
|
&router,
|
||||||
|
&client,
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
println!();
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute_prompt(
|
||||||
|
prompt: &str,
|
||||||
|
args: &ChatArgs,
|
||||||
|
config: &AppConfig,
|
||||||
|
budget_tracker: &mut BudgetTracker,
|
||||||
|
model_registry: &ModelRegistry,
|
||||||
|
router: &Router,
|
||||||
|
client: &OpenRouterClient,
|
||||||
|
) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
|
// 1. Budget check
|
||||||
|
match budget_tracker.check_budget() {
|
||||||
|
Ok(status) => {
|
||||||
|
if status.approaching_daily_limit {
|
||||||
|
eprintln!(
|
||||||
|
"{} Daily usage: {}/{}",
|
||||||
|
"⚠ Budget notice: Approaching daily ceiling.".yellow(),
|
||||||
|
status.daily_requests,
|
||||||
|
status.daily_limit
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!("{} {}", "Rate limit warning:".red().bold(), e);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let allow_stealth = args.allow_stealth || config.allow_stealth;
|
||||||
|
let require_paid = args.escalate;
|
||||||
|
|
||||||
|
// 2. Check if running a multi-agent pipeline
|
||||||
|
if let Some(ref pipe_name) = args.pipeline {
|
||||||
|
let pipeline = config
|
||||||
|
.pipelines
|
||||||
|
.iter()
|
||||||
|
.find(|p| p.name == *pipe_name)
|
||||||
|
.ok_or_else(|| format!("Pipeline '{pipe_name}' not found in configuration"))?;
|
||||||
|
|
||||||
|
run_pipeline(
|
||||||
|
pipeline,
|
||||||
|
config,
|
||||||
|
budget_tracker,
|
||||||
|
model_registry,
|
||||||
|
client,
|
||||||
|
prompt,
|
||||||
|
allow_stealth,
|
||||||
|
require_paid,
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. Single persona dispatch
|
||||||
|
let (target_persona, route_info) = if let Some(ref p) = args.persona {
|
||||||
|
(p.clone(), None)
|
||||||
|
} else {
|
||||||
|
let spinner = create_spinner("Routing message...");
|
||||||
|
// Record budget call if LLM router is invoked
|
||||||
|
let decision = router
|
||||||
|
.route(&config.persona_registry, client, prompt, require_paid)
|
||||||
|
.await;
|
||||||
|
spinner.finish_and_clear();
|
||||||
|
|
||||||
|
match decision {
|
||||||
|
Ok(dec) => {
|
||||||
|
if dec.stage_used == RouteStage::StageBLlm {
|
||||||
|
let _ = budget_tracker.record_call(false);
|
||||||
|
}
|
||||||
|
let info = match dec.stage_used {
|
||||||
|
RouteStage::StageAKeyword => "routed via keyword pre-filter".to_string(),
|
||||||
|
RouteStage::StageBLlm => "routed via LLM classifier".to_string(),
|
||||||
|
};
|
||||||
|
(dec.persona, Some(info))
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!("Router warning: {}, falling back to {}", e, config.default_persona);
|
||||||
|
(config.default_persona.clone(), None)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Some(info) = route_info {
|
||||||
|
eprintln!("{}", format!("⚡ [BrokeCrew: {info} ➔ {}]", target_persona).dimmed());
|
||||||
|
}
|
||||||
|
|
||||||
|
let (chunk_tx, mut chunk_rx) = mpsc::channel::<String>(100);
|
||||||
|
let persona_tag = get_persona_color_tag(&target_persona);
|
||||||
|
let spinner = create_spinner("Thinking...");
|
||||||
|
|
||||||
|
let dispatch_future = PipelineExecutor::dispatch_persona(
|
||||||
|
&target_persona,
|
||||||
|
&config.persona_registry,
|
||||||
|
model_registry,
|
||||||
|
client,
|
||||||
|
prompt,
|
||||||
|
allow_stealth,
|
||||||
|
require_paid,
|
||||||
|
Some(chunk_tx),
|
||||||
|
);
|
||||||
|
|
||||||
|
let printer = tokio::spawn(async move {
|
||||||
|
let mut first = true;
|
||||||
|
while let Some(chunk) = chunk_rx.recv().await {
|
||||||
|
if first {
|
||||||
|
spinner.finish_and_clear();
|
||||||
|
println!("\n{}", persona_tag);
|
||||||
|
first = false;
|
||||||
|
}
|
||||||
|
print!("{chunk}");
|
||||||
|
let _ = io::stdout().flush();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let dispatch_result = dispatch_future.await;
|
||||||
|
let _ = printer.await;
|
||||||
|
|
||||||
|
match dispatch_result {
|
||||||
|
Ok(res) => {
|
||||||
|
println!();
|
||||||
|
if res.fell_back {
|
||||||
|
eprintln!(
|
||||||
|
"{}",
|
||||||
|
format!("⚠ Candidate fallback: served by {}", res.model_used).yellow()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let _ = budget_tracker.record_call(res.tier_used == Tier::Paid);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!("\n{} {}", "Error dispatching to persona:".red().bold(), e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run_pipeline(
|
||||||
|
pipeline: &PipelineConfig,
|
||||||
|
config: &AppConfig,
|
||||||
|
budget_tracker: &mut BudgetTracker,
|
||||||
|
model_registry: &ModelRegistry,
|
||||||
|
client: &OpenRouterClient,
|
||||||
|
prompt: &str,
|
||||||
|
allow_stealth: bool,
|
||||||
|
require_paid: bool,
|
||||||
|
) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
|
println!(
|
||||||
|
"\n{} Starting pipeline '{}' ({} steps)",
|
||||||
|
"⚡".bright_cyan(),
|
||||||
|
pipeline.name.bold(),
|
||||||
|
pipeline.steps.len()
|
||||||
|
);
|
||||||
|
|
||||||
|
let (event_tx, mut event_rx) = mpsc::channel::<PipelineEvent>(100);
|
||||||
|
|
||||||
|
let exec_future = PipelineExecutor::execute(
|
||||||
|
pipeline,
|
||||||
|
&config.persona_registry,
|
||||||
|
model_registry,
|
||||||
|
client,
|
||||||
|
prompt,
|
||||||
|
allow_stealth,
|
||||||
|
require_paid,
|
||||||
|
Some(event_tx),
|
||||||
|
);
|
||||||
|
|
||||||
|
let event_handler = tokio::spawn(async move {
|
||||||
|
let mut current_spinner: Option<ProgressBar> = None;
|
||||||
|
|
||||||
|
while let Some(event) = event_rx.recv().await {
|
||||||
|
match event {
|
||||||
|
PipelineEvent::StepStarted {
|
||||||
|
step_index,
|
||||||
|
total_steps,
|
||||||
|
persona,
|
||||||
|
model,
|
||||||
|
..
|
||||||
|
} => {
|
||||||
|
if let Some(sp) = current_spinner.take() {
|
||||||
|
sp.finish_and_clear();
|
||||||
|
}
|
||||||
|
let badge = get_persona_color_tag(&persona);
|
||||||
|
println!(
|
||||||
|
"\n--- Step {}/{} {} ({}) ---",
|
||||||
|
step_index + 1,
|
||||||
|
total_steps,
|
||||||
|
badge,
|
||||||
|
model.dimmed()
|
||||||
|
);
|
||||||
|
let sp = create_spinner("Thinking...");
|
||||||
|
current_spinner = Some(sp);
|
||||||
|
}
|
||||||
|
PipelineEvent::Chunk { delta, .. } => {
|
||||||
|
if let Some(sp) = current_spinner.take() {
|
||||||
|
sp.finish_and_clear();
|
||||||
|
}
|
||||||
|
print!("{delta}");
|
||||||
|
let _ = io::stdout().flush();
|
||||||
|
}
|
||||||
|
PipelineEvent::StepCompleted { fell_back, .. } => {
|
||||||
|
if let Some(sp) = current_spinner.take() {
|
||||||
|
sp.finish_and_clear();
|
||||||
|
}
|
||||||
|
println!();
|
||||||
|
if fell_back {
|
||||||
|
eprintln!("{}", "⚠ Fallback occurred during step execution".yellow());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let results = exec_future.await?;
|
||||||
|
let _ = event_handler.await;
|
||||||
|
|
||||||
|
for r in &results {
|
||||||
|
let _ = budget_tracker.record_call(r.tier_used == Tier::Paid);
|
||||||
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue