wesichain_core/
runnable.rs1use 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 AwaitingApproval {
20 run_id: String,
22 prompt: String,
24 checkpoint_id: String,
26 },
27 ThinkingChunk(String),
29 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}