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 std::fmt;
11use std::sync::Arc;
12
13use ferrin_tool::ToolSet;
14use tokio_util::sync::CancellationToken;
15
16use super::result::EventStream;
17
18pub use smooth::ChunkDetector;
19pub use smooth::Chunking;
20pub use smooth::SmoothStream;
21pub use smooth::SmoothStreamConfig;
22pub use smooth::smooth_stream;
23
24/// Information handed to a transform when it is applied.
25#[derive(Clone)]
26pub struct TransformContext {
27    tools: Arc<ToolSet>,
28    cancellation: CancellationToken,
29    stop: CancellationToken,
30}
31
32impl TransformContext {
33    pub(crate) fn new(tools: Arc<ToolSet>, cancellation: CancellationToken) -> Self {
34        Self {
35            tools,
36            cancellation,
37            stop: CancellationToken::new(),
38        }
39    }
40
41    /// Token that fires when a transform calls [`stop`](Self::stop); the
42    /// pipeline gates further events on it.
43    pub(crate) fn stop_token(&self) -> CancellationToken {
44        self.stop.clone()
45    }
46
47    /// The tools available to the call.
48    #[must_use]
49    pub fn tools(&self) -> &ToolSet {
50        &self.tools
51    }
52
53    /// Stops the call: the model stream and pending tool executions are
54    /// cancelled and the call ends with [`crate::Error::Cancelled`]. A
55    /// transform that stops the call should also end its output stream so
56    /// that no further events reach the consumer.
57    pub fn stop(&self) {
58        self.stop.cancel();
59        self.cancellation.cancel();
60    }
61}
62
63impl fmt::Debug for TransformContext {
64    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
65        f.debug_struct("TransformContext")
66            .field("tools", &self.tools.names().collect::<Vec<_>>())
67            .field("stopped", &self.stop.is_cancelled())
68            .finish()
69    }
70}
71
72/// A stream transform.
73///
74/// Implemented for every `Fn(EventStream, TransformContext) -> EventStream`
75/// closure.
76pub trait StreamTransform: Send + Sync {
77    /// Wraps `input`.
78    fn apply(&self, input: EventStream, ctx: TransformContext) -> EventStream;
79}
80
81impl<F> StreamTransform for F
82where
83    F: Fn(EventStream, TransformContext) -> EventStream + Send + Sync,
84{
85    fn apply(&self, input: EventStream, ctx: TransformContext) -> EventStream {
86        self(input, ctx)
87    }
88}
ferrin_core/stream_text/transforms/mod.rs

ferrin_core/stream_text/transforms/
mod.rs
