427 lines
14 KiB
Rust
427 lines
14 KiB
Rust
use crate::client::{ClientError, ModelClient};
|
|
use crate::persona::{DispatchResult, ModelAttempt, ModelCandidate, 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<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct PipelineConfig {
|
|
pub name: String,
|
|
pub steps: Vec<PipelineStep>,
|
|
}
|
|
|
|
#[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<ModelAttempt>,
|
|
},
|
|
}
|
|
|
|
pub struct PipelineExecutor;
|
|
|
|
impl PipelineExecutor {
|
|
/// Executes a single persona query with candidate fallback walk
|
|
#[allow(clippy::too_many_arguments)]
|
|
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<mpsc::Sender<String>>,
|
|
) -> Result<DispatchResult, ClientError> {
|
|
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 dynamic_discovered;
|
|
let candidates_to_try: Vec<&ModelCandidate> = if !eligible.is_empty() {
|
|
eligible
|
|
} else if model_registry.model_count().await == 0 {
|
|
// Uninitialized / offline degraded state: allow raw candidates
|
|
persona
|
|
.candidates
|
|
.iter()
|
|
.filter(|c| allow_stealth || !c.is_stealth)
|
|
.collect()
|
|
} else {
|
|
// Live catalog was fetched, but declared candidates were rotated out.
|
|
// Dynamically discover active live free models for this persona:
|
|
dynamic_discovered = model_registry
|
|
.discover_free_candidates_for_role(persona_name)
|
|
.await;
|
|
if !dynamic_discovered.is_empty() {
|
|
tracing::info!(
|
|
persona = %persona_name,
|
|
count = dynamic_discovered.len(),
|
|
"Auto-discovered live free fallback candidates from OpenRouter catalog"
|
|
);
|
|
dynamic_discovered.iter().collect()
|
|
} else {
|
|
vec![]
|
|
}
|
|
};
|
|
|
|
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::<String>(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 {
|
|
let summary: Vec<String> = attempts
|
|
.iter()
|
|
.map(|a| {
|
|
format!(
|
|
"{}: {}",
|
|
a.model_id,
|
|
a.error.as_deref().unwrap_or("unknown error")
|
|
)
|
|
})
|
|
.collect();
|
|
return Err(ClientError::ModelUnavailable(format!(
|
|
"All {} candidates for persona '{persona_name}' failed:\n {}",
|
|
total_candidates,
|
|
summary.join("\n ")
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
let summary: Vec<String> = attempts
|
|
.iter()
|
|
.map(|a| {
|
|
format!(
|
|
"{}: {}",
|
|
a.model_id,
|
|
a.error.as_deref().unwrap_or("unknown error")
|
|
)
|
|
})
|
|
.collect();
|
|
Err(ClientError::ModelUnavailable(format!(
|
|
"All candidates for persona '{persona_name}' failed:\n {}",
|
|
summary.join("\n ")
|
|
)))
|
|
}
|
|
|
|
/// Executes an entire pipeline sequentially, passing outputs between steps
|
|
#[allow(clippy::too_many_arguments)]
|
|
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<mpsc::Sender<PipelineEvent>>,
|
|
) -> Result<Vec<DispatchResult>, 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::<String>(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<AtomicUsize>,
|
|
}
|
|
|
|
#[async_trait]
|
|
impl ModelClient for MockPipelineClient {
|
|
async fn complete(
|
|
&self,
|
|
model: &str,
|
|
_system: &str,
|
|
user: &str,
|
|
) -> Result<String, ClientError> {
|
|
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<String>,
|
|
) -> Result<String, ClientError> {
|
|
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);
|
|
}
|
|
}
|