wesichain-core 0.3.0

Rust-native LLM agents & chains with resumable ReAct workflows
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

use std::time::Duration;

use futures::stream::BoxStream;
use futures::StreamExt as _;

use crate::{Runnable, StreamEvent, WesichainError};

pub struct TimeLimited<R> {
    inner: R,
    timeout: Duration,
}

impl<R> TimeLimited<R> {
    pub fn new(inner: R, timeout: Duration) -> Self {
        Self { inner, timeout }
    }
}

#[async_trait::async_trait]
impl<Input, Output, R> Runnable<Input, Output> for TimeLimited<R>
where
    Input: Send + Clone + 'static,
    Output: Send + 'static,
    R: Runnable<Input, Output> + Send + Sync,
{
    async fn invoke(&self, input: Input) -> Result<Output, WesichainError> {
        tokio::time::timeout(self.timeout, self.inner.invoke(input))
            .await
            .map_err(|_| WesichainError::Timeout(self.timeout))?
    }

    fn stream<'a>(&'a self, input: Input) -> BoxStream<'a, Result<StreamEvent, WesichainError>> {
        let timeout = self.timeout;
        let inner = &self.inner;
        async_stream::stream! {
            let mut s = inner.stream(input);
            loop {
                match tokio::time::timeout(timeout, s.next()).await {
                    Err(_) => {
                        yield Err(WesichainError::Timeout(timeout));
                        break;
                    }
                    Ok(None) => break,
                    Ok(Some(event)) => yield event,
                }
            }
        }
        .boxed()
    }
}