ferrin_core/stream_text/transforms/
mod.rs

1//! Transforms applied to the event stream before step results are
2//! accumulated.
3//!
4//! A transform receives the whole event stream and returns a new one. It
5//! must preserve the event structure (start/delta/end pairs, step
6//! boundaries) so that the event processor can still build step results.
7
8mod smooth;
9
10use crate::error::Error;
11use std::fmt;
12use std::sync::Arc;
13use std::sync::Mutex;
14
15use ferrin_tool::ToolSet;
16use tokio_util::sync::CancellationToken;
17
18use super::result::EventStream;
19
20pub use smooth::ChunkDetector;
21pub use smooth::Chunking;
22pub use smooth::SmoothStream;
23pub use smooth::SmoothStreamConfig;
24pub use smooth::smooth_stream;
25
26/// Information handed to a transform when it is applied.
27#[derive(Clone)]
28pub struct TransformContext {
29    tools: Arc<ToolSet>,
30    cancellation: CancellationToken,
31    stop: CancellationToken,
32    failure: Arc<Mutex<Option<Error>>>,
33}
34
35impl TransformContext {
36    pub(crate) fn new(tools: Arc<ToolSet>, cancellation: CancellationToken) -> Self {
37        Self {
38            tools,
39            cancellation,
40            stop: CancellationToken::new(),
41            failure: Arc::new(Mutex::new(None)),
42        }
43    }
44
45    /// Token that fires when a transform calls [`stop`](Self::stop); the
46    /// pipeline gates further events on it.
47    pub(crate) fn stop_token(&self) -> CancellationToken {
48        self.stop.clone()
49    }
50
51    pub(crate) fn take_failure(&self) -> Option<Error> {
52        self.failure
53            .lock()
54            .unwrap_or_else(std::sync::PoisonError::into_inner)
55            .take()
56    }
57
58    pub(crate) fn cancellation(&self) -> &CancellationToken {
59        &self.cancellation
60    }
61
62    /// Stops the call with an application-supplied transform error.
63    ///
64    /// The first failure wins. End the transform's output stream after this
65    /// call; final-result waiters receive this error rather than cancellation.
66    pub fn fail(&self, error: Error) {
67        let mut failure = self
68            .failure
69            .lock()
70            .unwrap_or_else(std::sync::PoisonError::into_inner);
71        if failure.is_none() {
72            *failure = Some(error);
73        }
74        drop(failure);
75        self.stop();
76    }
77
78    /// The tools available to the call.
79    #[must_use]
80    pub fn tools(&self) -> &ToolSet {
81        &self.tools
82    }
83
84    /// Stops the call: the model stream and pending tool executions are
85    /// cancelled and the call ends with [`crate::Error::Cancelled`]. A
86    /// transform that stops the call should also end its output stream so
87    /// that no further events reach the consumer.
88    pub fn stop(&self) {
89        self.stop.cancel();
90        self.cancellation.cancel();
91    }
92}
93
94impl fmt::Debug for TransformContext {
95    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
96        f.debug_struct("TransformContext")
97            .field("tools", &self.tools.names().collect::<Vec<_>>())
98            .field("stopped", &self.stop.is_cancelled())
99            .finish()
100    }
101}
102
103/// A stream transform.
104///
105/// Implemented for every `Fn(EventStream, TransformContext) -> EventStream`
106/// closure.
107pub trait StreamTransform: Send + Sync {
108    /// Wraps `input`.
109    fn apply(&self, input: EventStream, ctx: TransformContext) -> EventStream;
110}
111
112impl<F> StreamTransform for F
113where
114    F: Fn(EventStream, TransformContext) -> EventStream + Send + Sync,
115{
116    fn apply(&self, input: EventStream, ctx: TransformContext) -> EventStream {
117        self(input, ctx)
118    }
119}
ferrin_core/stream_text/transforms/mod.rs

ferrin_core/stream_text/transforms/
mod.rs
