moly_kit/clients/
map.rs

1use crate::protocol::Tool;
2
3use crate::protocol::*;
4use crate::utils::asynchronous::{BoxPlatformSendFuture, BoxPlatformSendStream};
5use std::sync::{Arc, Mutex};
6
7struct Inner<C: BotClient> {
8    client: C,
9    map_bots: Option<Box<dyn FnMut(Vec<Bot>) -> Vec<Bot> + Send + 'static>>,
10    map_send: Option<Box<dyn FnMut(MessageContent) -> MessageContent + Send + 'static>>,
11}
12
13/// Utility wrapper client that transforms the output of the underlying client.
14///
15/// This is just limited to synchronous transformations over successful results.
16///
17/// In general, it's recommended to implement the [`BotClient`] trait directly for
18/// maximum control instead of using this.
19pub struct MapClient<C: BotClient> {
20    inner: Arc<Mutex<Inner<C>>>,
21}
22
23impl<C: BotClient> Clone for MapClient<C> {
24    fn clone(&self) -> Self {
25        MapClient {
26            inner: Arc::clone(&self.inner),
27        }
28    }
29}
30
31impl<C: BotClient> MapClient<C> {
32    pub fn new(client: C) -> Self {
33        MapClient {
34            inner: Arc::new(Mutex::new(Inner {
35                client,
36                map_bots: None,
37                map_send: None,
38            })),
39        }
40    }
41
42    /// Sets a transformation function for the list of bots returned by the `bots` method.
43    pub fn set_map_bots(&mut self, map: impl FnMut(Vec<Bot>) -> Vec<Bot> + Send + 'static) {
44        let mut inner = self.inner.lock().unwrap();
45        inner.map_bots = Some(Box::new(map));
46    }
47
48    /// Sets a sync transformation function for the successful result of the `send` method.
49    pub fn set_map_send(
50        &mut self,
51        map: impl FnMut(MessageContent) -> MessageContent + Send + 'static,
52    ) {
53        let mut inner = self.inner.lock().unwrap();
54        inner.map_send = Some(Box::new(map));
55    }
56}
57
58impl<C: BotClient> From<C> for MapClient<C> {
59    fn from(client: C) -> Self {
60        MapClient::new(client)
61    }
62}
63
64impl<C: BotClient + 'static> BotClient for MapClient<C> {
65    fn clone_box(&self) -> Box<dyn BotClient> {
66        Box::new(self.clone())
67    }
68
69    fn bots(&self) -> BoxPlatformSendFuture<'static, ClientResult<Vec<Bot>>> {
70        let inner = self.inner.clone();
71        let future = self.inner.lock().unwrap().client.bots();
72
73        Box::pin(async move {
74            let result = future.await;
75
76            if result.has_errors() {
77                return result;
78            }
79
80            let mut bots = result.into_value().unwrap();
81
82            if let Some(map_bots) = &mut inner.lock().unwrap().map_bots {
83                bots = map_bots(bots);
84            }
85
86            ClientResult::new_ok(bots)
87        })
88    }
89
90    fn send(
91        &mut self,
92        bot_id: &BotId,
93        messages: &[Message],
94        tools: &[Tool],
95    ) -> BoxPlatformSendStream<'static, ClientResult<MessageContent>> {
96        let inner = self.inner.clone();
97        let stream = self
98            .inner
99            .lock()
100            .unwrap()
101            .client
102            .send(bot_id, messages, tools);
103
104        let stream = async_stream::stream! {
105            for await result in stream {
106                if result.has_errors() {
107                    yield result;
108                    continue;
109                }
110
111                let mut content = result.into_value().unwrap();
112
113                if let Some(map_send) = &mut inner.lock().unwrap().map_send {
114                    content = map_send(content);
115                }
116
117                yield ClientResult::new_ok(content);
118            }
119        };
120
121        Box::pin(stream)
122    }
123}
moly_kit/clients/map.rs

moly_kit/clients/
map.rs
