gemini_adk_rs/text/
parallel.rs1use std::sync::Arc;
2
3use async_trait::async_trait;
4
5use super::TextAgent;
6use crate::context::AgentEvent;
7use crate::error::AgentError;
8use crate::middleware::MiddlewareChain;
9use crate::state::State;
10
11pub struct ParallelTextAgent {
14 name: String,
15 branches: Vec<Arc<dyn TextAgent>>,
16 middleware: MiddlewareChain,
17}
18
19impl ParallelTextAgent {
20 pub fn new(name: impl Into<String>, branches: Vec<Arc<dyn TextAgent>>) -> Self {
22 Self {
23 name: name.into(),
24 branches,
25 middleware: MiddlewareChain::new(),
26 }
27 }
28
29 pub fn with_middleware_chain(mut self, chain: MiddlewareChain) -> Self {
34 self.middleware = chain;
35 self
36 }
37}
38
39#[async_trait]
40impl TextAgent for ParallelTextAgent {
41 fn name(&self) -> &str {
42 &self.name
43 }
44
45 async fn run(&self, state: &State) -> Result<String, AgentError> {
46 let mut handles = Vec::with_capacity(self.branches.len());
47
48 for branch in &self.branches {
49 let _ = self
50 .middleware
51 .run_on_event(&AgentEvent::AgentStarted {
52 name: branch.name().to_string(),
53 })
54 .await;
55 let branch = branch.clone();
56 let state = state.clone();
57 handles.push(tokio::spawn(async move { branch.run(&state).await }));
58 }
59
60 let mut results = Vec::with_capacity(handles.len());
61 for (branch, handle) in self.branches.iter().zip(handles) {
62 let result = handle
63 .await
64 .map_err(|e| AgentError::Other(format!("Join error: {e}")))?;
65 results.push(result?);
66 let _ = self
67 .middleware
68 .run_on_event(&AgentEvent::AgentCompleted {
69 name: branch.name().to_string(),
70 })
71 .await;
72 }
73
74 let combined = results.join("\n");
75 let _ = state.set("output", &combined);
76 Ok(combined)
77 }
78}