From 271aeac23c443c980f31436f6cfdb49e144677da Mon Sep 17 00:00:00 2001 From: oiupoyt Date: Sat, 26 Sep 2026 13:50:04 +0530 Subject: [PATCH] feat: implement multi-agent handoff pipeline engine --- src/lib.rs | 2 + src/pipeline.rs | 377 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 379 insertions(+) create mode 100644 src/pipeline.rs diff --git a/src/lib.rs b/src/lib.rs index 6728615..87b277c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,11 +3,13 @@ pub mod budget; pub mod client; pub mod persona; +pub mod pipeline; pub mod registry; pub mod router; pub use budget::{BudgetExceeded, BudgetStatus, BudgetTracker, CreditTier}; pub use client::{ClientError, ModelClient, OpenRouterClient}; pub use persona::{DispatchResult, ModelAttempt, ModelCandidate, Persona, PersonaRegistry, Tier}; +pub use pipeline::{PipelineConfig, PipelineEvent, PipelineExecutor, PipelineStep}; pub use registry::{ModelInfo, ModelPricing, ModelRegistry}; pub use router::{KeywordRouter, LlmRouter, RouteDecision, RouteStage, Router}; diff --git a/src/pipeline.rs b/src/pipeline.rs new file mode 100644 index 0000000..2cca241 --- /dev/null +++ b/src/pipeline.rs @@ -0,0 +1,377 @@ +use crate::client::{ClientError, ModelClient}; +use crate::persona::{DispatchResult, ModelAttempt, PersonaRegistry, Tier}; +use crate::registry::ModelRegistry; +use serde::{Deserialize, Serialize}; +use tokio::sync::mpsc; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PipelineStep { + pub persona: String, + #[serde(default)] + pub pass_previous_output: bool, + #[serde(default)] + pub custom_prompt_template: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PipelineConfig { + pub name: String, + pub steps: Vec, +} + +#[derive(Debug, Clone)] +pub enum PipelineEvent { + StepStarted { + step_index: usize, + total_steps: usize, + persona: String, + model: String, + tier: Tier, + }, + Chunk { + step_index: usize, + delta: String, + }, + StepCompleted { + step_index: usize, + persona: String, + full_reply: String, + fell_back: bool, + attempts: Vec, + }, +} + +pub struct PipelineExecutor; + +impl PipelineExecutor { + /// Executes a single persona query with candidate fallback walk + pub async fn dispatch_persona( + persona_name: &str, + persona_registry: &PersonaRegistry, + model_registry: &ModelRegistry, + client: &dyn ModelClient, + prompt: &str, + allow_stealth: bool, + require_paid: bool, + chunk_sender: Option>, + ) -> Result { + let persona = persona_registry + .get(persona_name) + .ok_or_else(|| ClientError::ModelUnavailable(format!("Persona '{persona_name}' not found")))?; + + // Filter candidates against live registry and options + let eligible = model_registry + .filter_eligible(&persona.candidates, allow_stealth, require_paid) + .await; + + let candidates_to_try = if !eligible.is_empty() { + eligible + } else { + // If live registry filtering yielded nothing (e.g. offline/strict), fall back to raw persona candidates + persona.candidates.iter().filter(|c| allow_stealth || !c.is_stealth).collect() + }; + + if candidates_to_try.is_empty() { + return Err(ClientError::ModelUnavailable(format!( + "No eligible candidates available for persona '{persona_name}'" + ))); + } + + let mut attempts = Vec::new(); + let total_candidates = candidates_to_try.len(); + + for (idx, candidate) in candidates_to_try.iter().enumerate() { + let model_id = &candidate.id; + tracing::info!( + persona = %persona_name, + model = %model_id, + tier = ?candidate.tier, + attempt = idx + 1, + "Attempting candidate dispatch" + ); + + // Channel for streaming chunks if sender provided + let (tx, mut rx) = mpsc::channel::(100); + let forwarder = { + let external_sender = chunk_sender.clone(); + tokio::spawn(async move { + while let Some(chunk) = rx.recv().await { + if let Some(ref ext) = external_sender { + let _ = ext.send(chunk).await; + } + } + }) + }; + + let completion_result = client + .complete_stream(model_id, &persona.system_prompt, prompt, tx) + .await; + + let _ = forwarder.await; + + match completion_result { + Ok(reply) => { + let fell_back = idx > 0; + if fell_back { + tracing::warn!( + persona = %persona_name, + model = %model_id, + previous_failures = idx, + "Successful fallback to candidate" + ); + } + return Ok(DispatchResult { + persona: persona_name.to_string(), + model_used: model_id.clone(), + tier_used: candidate.tier, + reply, + fell_back, + attempts, + }); + } + Err(e) => { + let err_msg = e.to_string(); + tracing::warn!( + persona = %persona_name, + model = %model_id, + error = %err_msg, + "Candidate failed" + ); + + attempts.push(ModelAttempt { + model_id: model_id.clone(), + error: Some(err_msg.clone()), + }); + + // Fast-fail on auth or billing errors + match e { + ClientError::Unauthorized | ClientError::PaymentRequired => return Err(e), + _ => { + // Continue to next candidate + if idx + 1 == total_candidates { + return Err(ClientError::ModelUnavailable(format!( + "All {} candidates for persona '{persona_name}' failed. Last error: {err_msg}", + total_candidates + ))); + } + } + } + } + } + } + + Err(ClientError::ModelUnavailable(format!( + "All candidates for persona '{persona_name}' failed." + ))) + } + + /// Executes an entire pipeline sequentially, passing outputs between steps + pub async fn execute( + pipeline: &PipelineConfig, + persona_registry: &PersonaRegistry, + model_registry: &ModelRegistry, + client: &dyn ModelClient, + initial_prompt: &str, + allow_stealth: bool, + require_paid: bool, + event_sender: Option>, + ) -> Result, ClientError> { + let mut results = Vec::new(); + let mut previous_output = String::new(); + let total_steps = pipeline.steps.len(); + + for (step_idx, step) in pipeline.steps.iter().enumerate() { + let persona = persona_registry + .get(&step.persona) + .ok_or_else(|| ClientError::ModelUnavailable(format!("Persona '{}' not found", step.persona)))?; + + // Prepare prompt for this step + let step_prompt = if step.pass_previous_output && !previous_output.is_empty() { + if let Some(ref template) = step.custom_prompt_template { + template + .replace("{previous_output}", &previous_output) + .replace("{prompt}", initial_prompt) + } else { + format!( + "Context from previous step:\n```\n{}\n```\n\nTask:\n{}", + previous_output, initial_prompt + ) + } + } else { + initial_prompt.to_string() + }; + + // Setup chunk forwarding to PipelineEvent::Chunk + let (chunk_tx, mut chunk_rx) = mpsc::channel::(100); + let event_tx_clone = event_sender.clone(); + let forwarder = tokio::spawn(async move { + while let Some(chunk) = chunk_rx.recv().await { + if let Some(ref tx) = event_tx_clone { + let _ = tx + .send(PipelineEvent::Chunk { + step_index: step_idx, + delta: chunk, + }) + .await; + } + } + }); + + // Notify step start + if let Some(ref tx) = event_sender { + // Peek top candidate to inform UI + let top_candidate = persona + .candidates + .first() + .map(|c| (c.id.clone(), c.tier)) + .unwrap_or_else(|| ("unknown".into(), Tier::Free)); + + let _ = tx + .send(PipelineEvent::StepStarted { + step_index: step_idx, + total_steps, + persona: step.persona.clone(), + model: top_candidate.0, + tier: top_candidate.1, + }) + .await; + } + + let dispatch_result = Self::dispatch_persona( + &step.persona, + persona_registry, + model_registry, + client, + &step_prompt, + allow_stealth, + require_paid, + Some(chunk_tx), + ) + .await?; + + let _ = forwarder.await; + + previous_output = dispatch_result.reply.clone(); + + if let Some(ref tx) = event_sender { + let _ = tx + .send(PipelineEvent::StepCompleted { + step_index: step_idx, + persona: step.persona.clone(), + full_reply: dispatch_result.reply.clone(), + fell_back: dispatch_result.fell_back, + attempts: dispatch_result.attempts.clone(), + }) + .await; + } + + results.push(dispatch_result); + } + + Ok(results) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::persona::{ModelCandidate, Persona}; + use async_trait::async_trait; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + struct MockPipelineClient { + calls: Arc, + } + + #[async_trait] + impl ModelClient for MockPipelineClient { + async fn complete(&self, model: &str, _system: &str, user: &str) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(format!("Output from {} for: {}", model, user)) + } + + async fn complete_stream( + &self, + model: &str, + system: &str, + user: &str, + chunk_sender: mpsc::Sender, + ) -> Result { + let res = self.complete(model, system, user).await?; + let _ = chunk_sender.send(res.clone()).await; + Ok(res) + } + } + + #[tokio::test] + async fn test_pipeline_sequential_handoff() { + let coder = Persona { + name: "coder".into(), + system_prompt: "You are coder".into(), + trigger_words: vec![], + candidates: vec![ModelCandidate { + id: "coder-model".into(), + tier: Tier::Free, + is_stealth: false, + }], + }; + + let reviewer = Persona { + name: "reviewer".into(), + system_prompt: "You are reviewer".into(), + trigger_words: vec![], + candidates: vec![ModelCandidate { + id: "reviewer-model".into(), + tier: Tier::Free, + is_stealth: false, + }], + }; + + let persona_reg = PersonaRegistry::new(vec![coder, reviewer]); + let model_reg = ModelRegistry::new(); + + let pipeline = PipelineConfig { + name: "build-and-review".into(), + steps: vec![ + PipelineStep { + persona: "coder".into(), + pass_previous_output: false, + custom_prompt_template: None, + }, + PipelineStep { + persona: "reviewer".into(), + pass_previous_output: true, + custom_prompt_template: None, + }, + ], + }; + + let calls = Arc::new(AtomicUsize::new(0)); + let client = MockPipelineClient { + calls: calls.clone(), + }; + + let results = PipelineExecutor::execute( + &pipeline, + &persona_reg, + &model_reg, + &client, + "Build website", + false, + false, + None, + ) + .await + .expect("pipeline execution failed"); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].persona, "coder"); + assert_eq!(results[0].model_used, "coder-model"); + assert_eq!(results[1].persona, "reviewer"); + assert_eq!(results[1].model_used, "reviewer-model"); + // Reviewer step should contain output from coder step + assert!(results[1].reply.contains("Output from coder-model")); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } +}