vibe-action 0.2.2

Command router — execute shell commands and LLM prompts via simple YAML actions.
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

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

use std::time::Duration;

use anyhow::Result;
use vibe_cluster::BatchOptions;
use vibe_cluster::Prompt;

use crate::configs::app::AppConfig;
use crate::engine::parser;
use crate::models::action::ActionRun;
use crate::output::output::OutputKind;
use crate::print_template;
use crate::utils;

/// 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],
        images: Option<Vec<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,
                user: parser::unescape_text(p),
                system: Some(system.to_string()),
                images: images.clone(),
            })
            .collect();

        let cluster_results = cluster
            .batch_call_with_options(&cluster_prompts, &options, move |result, current, total| {
                // Get size image
                let image_info = result
                    .prompt
                    .images
                    .as_ref()
                    .map(|imgs| {
                        let size: usize = imgs.iter().map(|i| i.len()).sum();
                        format!(", {} image", utils::format::format_bytes(size))
                    })
                    .unwrap_or_default();

                // Get size token
                let bpe = tiktoken_rs::o200k_base().unwrap();
                // Approximate token count for non-OpenAI models
                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);

                // Get label
                let role_label = AppConfig::instance()
                    .ok()
                    .and_then(|c| c.cluster.iter().find(|n| n.model == result.model))
                    .and_then(|n| n.role.as_ref())
                    .map(|r| r.to_string())
                    .unwrap_or_else(|| "llm".to_string());

                // Get duration
                let duration_ms =
                    utils::format::format_duration(Duration::from_millis(result.duration_ms));

                // Print result info
                print_template!(
                    OutputKind::Progress,
                    "[{role}] node batch: {current}/{total} | {model} | {duration} (~{tokens} tokens{image})",
                    "role" => role_label,
                    "current" => format!("{:>width$}", current, width = total.to_string().len()),
                    "total" => total.to_string(),
                    "model" => result.model,
                    "duration" => duration_ms,
                    "tokens" => (user_tokens + system_tokens).to_string(),
                    "image" => image_info,
                );
            })
            .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 {
                // @todo
                // println!("----");
                // println!("{}", &text);
                // println!("----");
                final_results.push(ClusterResult {
                    prompt: result.prompt.user,
                    result: vibe_cluster::normalize_text(&text),
                });
            }
        }

        Ok(final_results)
    }
}