Skip to main content

wesichain_core/
runnable.rs

1use async_trait::async_trait;
2use futures::stream::BoxStream;
3use std::sync::Arc;
4
5use crate::{serde::SerializableRunnable, WesichainError};
6
7#[derive(Debug, Clone, PartialEq)]
8pub enum StreamEvent {
9    ContentChunk(String),
10    ToolCallStart { id: String, name: String },
11    ToolCallDelta { id: String, delta: crate::Value },
12    ToolCallResult { id: String, output: crate::Value },
13    FinalAnswer(String),
14    Metadata { key: String, value: crate::Value },
15    /// Emitted when an agent pauses at a human-in-the-loop approval gate.
16    ///
17    /// The host application should present `prompt` to a human, collect their
18    /// decision, and resume the graph from `checkpoint_id`.
19    AwaitingApproval {
20        /// Stable identifier for this execution run.
21        run_id: String,
22        /// Human-readable description of what the agent wants to do.
23        prompt: String,
24        /// Opaque checkpoint identifier — pass back to resume the graph.
25        checkpoint_id: String,
26    },
27    /// Model reasoning / scratchpad text (Anthropic extended thinking, OpenAI o-series).
28    ThinkingChunk(String),
29    /// Token consumption update emitted after each LLM response.
30    UsageUpdate {
31        input_tokens: u32,
32        output_tokens: u32,
33        cache_read_tokens: Option<u32>,
34        cache_write_tokens: Option<u32>,
35    },
36}
37
38#[async_trait]
39pub trait Runnable<Input: Send + 'static, Output: Send + 'static>: Send + Sync {
40    async fn invoke(&self, input: Input) -> Result<Output, WesichainError>;
41
42    async fn batch(&self, inputs: Vec<Input>) -> Vec<Result<Output, WesichainError>> {
43        let futures = inputs.into_iter().map(|i| self.invoke(i));
44        futures::future::join_all(futures).await
45    }
46
47    async fn abatch(&self, inputs: Vec<Input>) -> Vec<Result<Output, WesichainError>> {
48        self.batch(inputs).await
49    }
50
51    fn to_serializable(&self) -> Option<SerializableRunnable> {
52        None
53    }
54
55    fn stream<'a>(&'a self, input: Input) -> BoxStream<'a, Result<StreamEvent, WesichainError>>;
56}
57
58#[async_trait]
59impl<Input, Output, T> Runnable<Input, Output> for Arc<T>
60where
61    Input: Send + 'static,
62    Output: Send + 'static,
63    T: Runnable<Input, Output> + ?Sized,
64{
65    async fn invoke(&self, input: Input) -> Result<Output, WesichainError> {
66        (**self).invoke(input).await
67    }
68
69    fn stream<'a>(&'a self, input: Input) -> BoxStream<'a, Result<StreamEvent, WesichainError>> {
70        (**self).stream(input)
71    }
72
73    fn to_serializable(&self) -> Option<SerializableRunnable> {
74        (**self).to_serializable()
75    }
76}
77
78#[async_trait]
79impl<Input, Output, T> Runnable<Input, Output> for Box<T>
80where
81    Input: Send + 'static,
82    Output: Send + 'static,
83    T: Runnable<Input, Output> + ?Sized,
84{
85    async fn invoke(&self, input: Input) -> Result<Output, WesichainError> {
86        (**self).invoke(input).await
87    }
88
89    fn stream<'a>(&'a self, input: Input) -> BoxStream<'a, Result<StreamEvent, WesichainError>> {
90        (**self).stream(input)
91    }
92
93    fn to_serializable(&self) -> Option<SerializableRunnable> {
94        (**self).to_serializable()
95    }
96}