vibe-action 0.0.5

//! LLM cluster executor.
//! Sends prompts to vibe-cluster and returns model responses.

use anyhow::Result;
use vibe_cluster::{BatchOptions, Prompt};

use crate::{
    configs::app::AppConfig, models::action::ActionRun, output::output::OutputLevel, print_info,
    print_progress,
};

/// Distributed cluster execution response payload.
pub struct ClusterResult {
    /// The generated text prompt or command payload to be executed.
    pub prompt: String,
    /// Normalized text response generated by the cluster node model.
    pub result: String,
}

pub struct Cluster;

impl Cluster {
    /// Execute an LLM prompt batch concurrently across the distributed cluster.
    pub async fn exec(
        system: &str,
        retries: u32,
        run: &ActionRun,
        prompts: &[String],
    ) -> Result<Vec<ClusterResult>> {
        let config = AppConfig::instance()?;
        let cluster = config.create_cluster_filtered(run)?;

        let options = BatchOptions {
            retries: Some(retries as usize),
            ..BatchOptions::default()
        };

        let cluster_prompts: Vec<Prompt> = prompts
            .iter()
            .map(|p| Prompt {
                key: None,
                system: Some(system.to_string()),
                user: p.clone(),
            })
            .collect();

        let run_arc = std::sync::Arc::new(run.to_string());
        let run_for_closure = run_arc.clone();

        let cluster_results = cluster
            .batch_call_with_options(&cluster_prompts, &options, move |result, current, total| {
                let bpe = tiktoken_rs::cl100k_base().unwrap();
                let user_tokens = bpe.encode_with_special_tokens(&result.prompt.user).len();
                let system_tokens = result
                    .prompt
                    .system
                    .as_ref()
                    .map(|s| bpe.encode_with_special_tokens(s).len())
                    .unwrap_or(0);
                let total_width = total.to_string().len();
                let total_tokens = user_tokens + system_tokens;
                if AppConfig::output().level() == OutputLevel::Cli {
                    print_progress!(
                        "└─ [{}] node batch: {:>width$}/{} | {} finished in {}ms ({} tokens)",
                        run_for_closure,
                        current,
                        total,
                        result.model,
                        result.duration_ms,
                        total_tokens,
                        width = total_width
                    );
                } else {
                    print_info!(
                        "[{}] batch {:>width$}/{} completed by node '{}' in {}ms ({} tokens)",
                        run_for_closure,
                        current,
                        total,
                        result.model,
                        result.duration_ms,
                        total_tokens,
                        width = total_width
                    );
                }
            })
            .await
            .map_err(|e| anyhow::anyhow!("Cluster batch call failed: {}", e))?;

        let mut final_results = Vec::with_capacity(cluster_results.len());

        for result in cluster_results {
            if let Some(err) = result.error {
                anyhow::bail!("Cluster node execution failed: {}", err);
            }
            if let Some(text) = result.text {
                final_results.push(ClusterResult {
                    prompt: result.prompt.user,
                    result: vibe_cluster::normalize_text(&text),
                });
            }
        }

        Ok(final_results)
    }
}