stasis/application/orchestration/
concurrent_pattern_pipeline.rs

1use std::sync::Arc;
2
3use crate::application::orchestration::prompt_pipeline::{
4    PromptExecutionContext, PromptExecutionPipeline, PromptExecutionRequest,
5};
6use crate::domain::errors::{Result, StasisError};
7use tokio::task::JoinSet;
8
9#[derive(Clone, Debug)]
10pub struct ConcurrentPatternBranch {
11    pub branch_id: String,
12    pub user_prompt_template: String,
13    pub system_prompt: Option<String>,
14    pub policy_profile: Option<String>,
15    pub model_hint: Option<String>,
16}
17
18#[derive(Clone, Debug)]
19pub struct ConcurrentPatternExecutionRequest {
20    pub initial_user_prompt: String,
21    pub trace_id: Option<String>,
22    pub correlation_id: Option<String>,
23    pub policy_profile: Option<String>,
24    pub model_hint: Option<String>,
25    pub merge_strategy: Option<String>,
26    pub branches: Vec<ConcurrentPatternBranch>,
27}
28
29#[derive(Clone, Debug)]
30pub struct ConcurrentPatternBranchResult {
31    pub branch_id: String,
32    pub rendered_prompt: String,
33    pub output_text: String,
34}
35
36#[derive(Clone, Debug)]
37pub struct ConcurrentPatternExecutionResponse {
38    pub final_text: String,
39    pub branches: Vec<ConcurrentPatternBranchResult>,
40    pub termination_reason: String,
41    pub merge_strategy: String,
42}
43
44#[derive(Clone)]
45pub struct ConcurrentPatternPipeline {
46    prompt_pipeline: PromptExecutionPipeline,
47}
48
49#[derive(Clone)]
50struct ConcurrentSharedInputs {
51    initial_input: Arc<str>,
52    trace_id: Arc<Option<String>>,
53    correlation_id: Arc<Option<String>>,
54    default_policy_profile: Arc<Option<String>>,
55    default_model_hint: Arc<Option<String>>,
56}
57
58impl ConcurrentSharedInputs {
59    fn build_context(
60        &self,
61        policy_profile: Option<String>,
62        model_hint: Option<String>,
63    ) -> PromptExecutionContext {
64        PromptExecutionContext {
65            trace_id: (*self.trace_id).clone(),
66            correlation_id: (*self.correlation_id).clone(),
67            policy_profile: policy_profile.or_else(|| (*self.default_policy_profile).clone()),
68            model_hint: model_hint.or_else(|| (*self.default_model_hint).clone()),
69        }
70    }
71}
72
73impl ConcurrentPatternPipeline {
74    pub fn new(prompt_pipeline: PromptExecutionPipeline) -> Self {
75        Self { prompt_pipeline }
76    }
77
78    pub async fn execute(
79        &self,
80        request: ConcurrentPatternExecutionRequest,
81    ) -> Result<ConcurrentPatternExecutionResponse> {
82        let ConcurrentPatternExecutionRequest {
83            initial_user_prompt,
84            trace_id,
85            correlation_id,
86            policy_profile,
87            model_hint,
88            merge_strategy,
89            branches,
90        } = request;
91
92        let merge_strategy = merge_strategy.unwrap_or_else(|| "join_with_headers".to_string());
93        let shared_inputs = ConcurrentSharedInputs {
94            initial_input: Arc::<str>::from(initial_user_prompt),
95            trace_id: Arc::new(trace_id),
96            correlation_id: Arc::new(correlation_id),
97            default_policy_profile: Arc::new(policy_profile),
98            default_model_hint: Arc::new(model_hint),
99        };
100
101        let mut join_set: JoinSet<Result<ConcurrentPatternBranchResult>> = JoinSet::new();
102
103        for branch in branches {
104            let pipeline = self.prompt_pipeline.clone();
105            let shared_inputs = shared_inputs.clone();
106
107            join_set.spawn(async move {
108                let ConcurrentPatternBranch {
109                    branch_id,
110                    user_prompt_template,
111                    system_prompt,
112                    policy_profile,
113                    model_hint,
114                } = branch;
115
116                let rendered_prompt = user_prompt_template
117                    .replace("{{input}}", &shared_inputs.initial_input)
118                    .replace("{input}", &shared_inputs.initial_input);
119
120                let context = shared_inputs.build_context(policy_profile, model_hint);
121
122                let mut prompt_request =
123                    PromptExecutionRequest::from_user_prompt(rendered_prompt.clone())
124                        .with_context(context);
125                if let Some(system_prompt) = system_prompt {
126                    prompt_request = prompt_request.with_system_prompt(system_prompt);
127                }
128
129                let response = pipeline.execute(prompt_request).await?;
130                Ok(ConcurrentPatternBranchResult {
131                    branch_id,
132                    rendered_prompt,
133                    output_text: response.text,
134                })
135            });
136        }
137
138        let mut results = Vec::new();
139        while let Some(joined) = join_set.join_next().await {
140            let result = joined.map_err(|err| {
141                StasisError::PortFailure(format!("concurrent pattern join failure: {err}"))
142            })??;
143            results.push(result);
144        }
145
146        results.sort_by(|a, b| a.branch_id.cmp(&b.branch_id));
147
148        let final_text = render_final_text(&results, &merge_strategy);
149
150        Ok(ConcurrentPatternExecutionResponse {
151            final_text,
152            branches: results,
153            termination_reason: "completed_all_branches".to_string(),
154            merge_strategy,
155        })
156    }
157}
158
159fn render_final_text(results: &[ConcurrentPatternBranchResult], merge_strategy: &str) -> String {
160    if results.is_empty() {
161        return String::new();
162    }
163
164    match merge_strategy {
165        "join_lines" => {
166            let total_text_len: usize = results.iter().map(|branch| branch.output_text.len()).sum();
167            let mut final_text = String::with_capacity(total_text_len + results.len().saturating_sub(1));
168            for (idx, branch) in results.iter().enumerate() {
169                if idx > 0 {
170                    final_text.push('\n');
171                }
172                final_text.push_str(&branch.output_text);
173            }
174            final_text
175        }
176        _ => {
177            let total_text_len: usize = results
178                .iter()
179                .map(|branch| branch.branch_id.len() + branch.output_text.len() + 4)
180                .sum();
181            let separator_len = 2 * results.len().saturating_sub(1);
182            let mut final_text = String::with_capacity(total_text_len + separator_len);
183            for (idx, branch) in results.iter().enumerate() {
184                if idx > 0 {
185                    final_text.push_str("\n\n");
186                }
187                final_text.push('[');
188                final_text.push_str(&branch.branch_id);
189                final_text.push_str("]\n");
190                final_text.push_str(&branch.output_text);
191            }
192            final_text
193        }
194    }
195}
stasis/application/orchestration/concurrent_pattern_pipeline.rs

stasis/application/orchestration/
concurrent_pattern_pipeline.rs
