moly_kit/clients/
multi.rs

1use crate::protocol::Tool;
2use makepad_widgets::{Cx, LiveId, LivePtr, WidgetRef};
3
4use crate::protocol::*;
5use crate::utils::asynchronous::{BoxPlatformSendFuture, BoxPlatformSendStream};
6use std::{
7    collections::HashMap,
8    sync::{Arc, Mutex},
9};
10
11/// A client that can be composed from multiple subclients to interact with all of them as one.
12#[derive(Clone)]
13pub struct MultiClient {
14    clients_with_bots: Arc<Mutex<Vec<(Box<dyn BotClient>, Vec<Bot>)>>>,
15}
16
17impl MultiClient {
18    pub fn new() -> Self {
19        MultiClient {
20            clients_with_bots: Arc::new(Mutex::new(Vec::new())),
21        }
22    }
23
24    pub fn add_client(&mut self, client: Box<dyn BotClient>) {
25        self.clients_with_bots
26            .lock()
27            .unwrap()
28            .push((client, Vec::new()));
29    }
30}
31
32impl BotClient for MultiClient {
33    fn send(
34        &mut self,
35        bot_id: &BotId,
36        messages: &[Message],
37        tools: &[Tool],
38    ) -> BoxPlatformSendStream<'static, ClientResult<MessageContent>> {
39        let client = self
40            .clients_with_bots
41            .lock()
42            .unwrap()
43            .iter()
44            .find_map(|(client, bots)| {
45                if bots.iter().any(|b| b.id == *bot_id) {
46                    Some(client.clone())
47                } else {
48                    None
49                }
50            });
51
52        match client {
53            Some(mut client) => client.send(bot_id, messages, tools),
54            None => {
55                let bot_id_clone = bot_id.clone();
56                Box::pin(futures::stream::once(async move {
57                    ClientError::new(
58                        ClientErrorKind::Unknown,
59                        format!(
60                            "Can't find a client to communicate with the bot {:?}",
61                            bot_id_clone
62                        ),
63                    )
64                    .into()
65                }))
66            }
67        }
68    }
69
70    // TODO: Add `send` implementation to take adventage of `send` implementation in sub-clients.
71
72    fn clone_box(&self) -> Box<dyn BotClient> {
73        Box::new(self.clone())
74    }
75
76    fn bots(&self) -> BoxPlatformSendFuture<'static, ClientResult<Vec<Bot>>> {
77        let clients_with_bots = self.clients_with_bots.clone();
78
79        let future = async move {
80            let clients = clients_with_bots
81                .lock()
82                .unwrap()
83                .iter()
84                .map(|(client, _)| client.clone())
85                .collect::<Vec<_>>();
86
87            let bot_futures = clients.iter().map(|client| client.bots());
88            let results = futures::future::join_all(bot_futures).await;
89
90            let mut zipped_bots = Vec::new();
91            let mut flat_bots = Vec::new();
92            let mut errors = Vec::new();
93
94            for result in results {
95                let (v, e) = result.into_value_and_errors();
96                let v = v.unwrap_or_default();
97                zipped_bots.push(v.clone());
98                flat_bots.extend(v);
99                errors.extend(e);
100            }
101
102            *clients_with_bots.lock().unwrap() = clients
103                .into_iter()
104                .zip(zipped_bots.iter().cloned())
105                .collect();
106
107            if errors.is_empty() {
108                ClientResult::new_ok(flat_bots)
109            } else {
110                if flat_bots.is_empty() {
111                    ClientResult::new_err(errors)
112                } else {
113                    ClientResult::new_ok_and_err(flat_bots, errors)
114                }
115            }
116        };
117
118        Box::pin(future)
119    }
120
121    fn content_widget(
122        &mut self,
123        cx: &mut Cx,
124        previous_widget: WidgetRef,
125        templates: &HashMap<LiveId, LivePtr>,
126        content: &MessageContent,
127    ) -> Option<WidgetRef> {
128        self.clients_with_bots
129            .lock()
130            .unwrap()
131            .iter_mut()
132            .find_map(|(client, _)| {
133                client.content_widget(cx, previous_widget.clone(), templates, content)
134            })
135    }
136}