ferrin_core/stream_text/
builder.rs

1//! The `stream_text` builder.
2
3use std::fmt;
4use std::future::Future;
5use std::future::IntoFuture;
6use std::sync::Arc;
7
8use ferrin_spec::BoxFuture;
9use ferrin_spec::LanguageModelRef;
10
11use super::StreamErrorInfo;
12use super::StreamEvent;
13use super::pipeline;
14use super::result::StreamTextResult;
15use super::transforms::StreamTransform;
16use crate::error::Error;
17use crate::generate_text::config::CallConfig;
18use crate::hooks::HookFn;
19use crate::output::NoOutput;
20use crate::output::Output;
21use crate::output::OutputHandler;
22use crate::telemetry::AbortEvent;
23
24/// What the pipeline does after an error reported by the model stream.
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
26#[non_exhaustive]
27pub enum ErrorDecision {
28    /// Report the error and end the call.
29    #[default]
30    Continue,
31    /// Retry the model call of the current step. Honoured at most once per
32    /// step and only when [`StreamText::stream_retries`] is configured.
33    Retry,
34}
35
36/// Callback invoked for every error that occurs while streaming.
37///
38/// Implemented for every `Fn(StreamErrorInfo) -> impl Future<Output =
39/// ErrorDecision>` closure.
40pub trait OnErrorFn: Send + Sync + 'static {
41    /// Handles one error and decides how to proceed.
42    fn call(&self, error: StreamErrorInfo) -> BoxFuture<'static, ErrorDecision>;
43}
44
45impl<F, Fut> OnErrorFn for F
46where
47    F: Fn(StreamErrorInfo) -> Fut + Send + Sync + 'static,
48    Fut: Future<Output = ErrorDecision> + Send + 'static,
49{
50    fn call(&self, error: StreamErrorInfo) -> BoxFuture<'static, ErrorDecision> {
51        Box::pin(self(error))
52    }
53}
54
55/// Streaming-only settings.
56#[derive(Default)]
57pub(crate) struct StreamConfig {
58    pub(crate) transforms: Vec<Arc<dyn StreamTransform>>,
59    pub(crate) include_raw_chunks: bool,
60    pub(crate) stream_retries: Option<u32>,
61    pub(crate) on_error: Option<Arc<dyn OnErrorFn>>,
62}
63
64impl fmt::Debug for StreamConfig {
65    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
66        f.debug_struct("StreamConfig")
67            .field("transforms", &self.transforms.len())
68            .field("include_raw_chunks", &self.include_raw_chunks)
69            .field("stream_retries", &self.stream_retries)
70            .field("on_error", &self.on_error.is_some())
71            .finish()
72    }
73}
74
75/// Starts building a streaming text generation call.
76///
77/// The builder is a future: `.await` starts the pipeline and resolves once
78/// the first model request has been established (or failed after retries).
79#[must_use]
80pub fn stream_text(model: impl Into<LanguageModelRef>) -> StreamText<()> {
81    StreamText {
82        config: CallConfig::new(model.into()),
83        output: Arc::new(NoOutput),
84        stream: StreamConfig {
85            transforms: Vec::new(),
86            include_raw_chunks: false,
87            stream_retries: None,
88            on_error: None,
89        },
90    }
91}
92
93/// Builder and future of a `stream_text` call.
94pub struct StreamText<O> {
95    pub(crate) config: CallConfig,
96    pub(crate) output: Arc<dyn OutputHandler<O>>,
97    pub(crate) stream: StreamConfig,
98}
99
100crate::builder::impl_call_builder!(StreamText);
101
102impl<O> StreamText<O> {
103    /// Requests structured output parsed by `output`.
104    pub fn output<T>(self, output: Output<T>) -> StreamText<T> {
105        StreamText {
106            config: self.config,
107            output: output.handler(),
108            stream: self.stream,
109        }
110    }
111
112    /// Appends a transform applied to the event stream before the step
113    /// results are accumulated. Transforms run in order.
114    #[must_use]
115    pub fn transform(mut self, transform: impl StreamTransform + 'static) -> Self {
116        self.stream.transforms.push(Arc::new(transform));
117        self
118    }
119
120    /// Forwards raw provider chunks as [`StreamEvent::Raw`].
121    #[must_use]
122    pub fn include_raw_chunks(mut self) -> Self {
123        self.stream.include_raw_chunks = true;
124        self
125    }
126
127    /// Enables stream-level retries: after an error reported by the model
128    /// stream the current step is retried up to `retries` times. `0` allows
129    /// only retries requested by the [`on_error`](Self::on_error) callback.
130    /// Retries are disabled when this is not called.
131    #[must_use]
132    pub fn stream_retries(mut self, retries: u32) -> Self {
133        self.stream.stream_retries = Some(retries);
134        self
135    }
136
137    /// Runs for every emitted event (after transforms).
138    #[must_use]
139    pub fn on_chunk(mut self, f: impl HookFn<StreamEvent>) -> Self {
140        self.config.hooks.on_chunk.push(Arc::new(f));
141        self
142    }
143
144    /// Runs when the call is aborted by cancellation.
145    #[must_use]
146    pub fn on_abort(mut self, f: impl HookFn<AbortEvent>) -> Self {
147        self.config.hooks.on_abort.push(Arc::new(f));
148        self
149    }
150
151    /// Runs for every error that occurs while streaming; the returned
152    /// decision may request a retry (see [`stream_retries`](Self::stream_retries)).
153    #[must_use]
154    pub fn on_error(mut self, f: impl OnErrorFn) -> Self {
155        self.stream.on_error = Some(Arc::new(f));
156        self
157    }
158}
159
160impl<O> fmt::Debug for StreamText<O> {
161    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
162        f.debug_struct("StreamText")
163            .field("config", &self.config)
164            .field("stream", &self.stream)
165            .finish_non_exhaustive()
166    }
167}
168
169impl<O: Send + 'static> IntoFuture for StreamText<O> {
170    type Output = Result<StreamTextResult<O>, Error>;
171    type IntoFuture = BoxFuture<'static, Self::Output>;
172
173    fn into_future(self) -> Self::IntoFuture {
174        Box::pin(pipeline::start(self.config, self.output, self.stream))
175    }
176}
ferrin_core/stream_text/builder.rs

ferrin_core/stream_text/
builder.rs
