Skip to main content

ferrin_core/middleware/builtin/
simulate_streaming.rs

1//! Streaming simulated from a complete response.
2
3use ferrin_spec::BoxFuture;
4use ferrin_spec::CallOptions;
5use ferrin_spec::Content;
6use ferrin_spec::PartId;
7use ferrin_spec::StreamPart;
8use ferrin_spec::StreamResult;
9use ferrin_spec::error::ProviderError;
10use ferrin_spec::language_model::GenerateResult;
11use futures_util::stream;
12
13use crate::middleware::LanguageModelMiddleware;
14use crate::middleware::MiddlewareContext;
15use crate::middleware::StreamNext;
16
17/// Middleware created by [`simulate_streaming`].
18#[derive(Debug, Clone, Copy, Default)]
19pub struct SimulateStreaming;
20
21/// Serves `do_stream` by calling `do_generate` on the wrapped model and
22/// expanding the result into `stream-start`, `response-metadata`, one
23/// start/delta/end triple per text or reasoning part (empty text parts are
24/// skipped), the other parts as is, and `finish`.
25#[must_use]
26pub fn simulate_streaming() -> SimulateStreaming {
27    SimulateStreaming
28}
29
30/// Expands a complete result into stream parts.
31#[must_use]
32pub fn simulate_parts(result: &GenerateResult) -> Vec<StreamPart> {
33    let mut parts = vec![
34        StreamPart::StreamStart {
35            warnings: result.warnings.clone(),
36        },
37        StreamPart::ResponseMetadata {
38            id: result.response.id.clone(),
39            timestamp: result.response.timestamp,
40            model_id: result.response.model_id.clone(),
41        },
42    ];
43    let mut next_id = 0u32;
44    for part in &result.content {
45        match part {
46            Content::Text {
47                text,
48                provider_metadata,
49            } => {
50                if text.is_empty() {
51                    continue;
52                }
53                let id = PartId::new(next_id.to_string());
54                next_id += 1;
55                parts.push(StreamPart::TextStart {
56                    id: id.clone(),
57                    provider_metadata: provider_metadata.clone(),
58                });
59                parts.push(StreamPart::TextDelta {
60                    id: id.clone(),
61                    delta: text.clone(),
62                    provider_metadata: None,
63                });
64                parts.push(StreamPart::TextEnd {
65                    id,
66                    provider_metadata: None,
67                });
68            }
69            Content::Reasoning {
70                text,
71                provider_metadata,
72            } => {
73                let id = PartId::new(next_id.to_string());
74                next_id += 1;
75                parts.push(StreamPart::ReasoningStart {
76                    id: id.clone(),
77                    provider_metadata: provider_metadata.clone(),
78                });
79                parts.push(StreamPart::ReasoningDelta {
80                    id: id.clone(),
81                    delta: text.clone(),
82                    provider_metadata: None,
83                });
84                parts.push(StreamPart::ReasoningEnd {
85                    id,
86                    provider_metadata: None,
87                });
88            }
89            Content::ReasoningFile {
90                data,
91                media_type,
92                provider_metadata,
93            } => parts.push(StreamPart::ReasoningFile {
94                data: data.clone(),
95                media_type: media_type.clone(),
96                provider_metadata: provider_metadata.clone(),
97            }),
98            Content::File {
99                data,
100                media_type,
101                filename,
102                provider_metadata,
103            } => parts.push(StreamPart::File {
104                data: data.clone(),
105                media_type: media_type.clone(),
106                filename: filename.clone(),
107                provider_metadata: provider_metadata.clone(),
108            }),
109            Content::Custom {
110                kind,
111                provider_metadata,
112            } => parts.push(StreamPart::Custom {
113                kind: kind.clone(),
114                provider_metadata: provider_metadata.clone(),
115            }),
116            Content::Source(source) => parts.push(StreamPart::Source(source.clone())),
117            Content::ToolCall(call) => parts.push(StreamPart::ToolCall(call.clone())),
118            Content::ToolResult(result) => parts.push(StreamPart::ToolResult(result.clone())),
119            Content::ToolApprovalRequest {
120                approval_id,
121                tool_call_id,
122                provider_metadata,
123            } => parts.push(StreamPart::ToolApprovalRequest {
124                approval_id: approval_id.clone(),
125                tool_call_id: tool_call_id.clone(),
126                provider_metadata: provider_metadata.clone(),
127            }),
128            #[allow(unreachable_patterns, reason = "Content is non-exhaustive")]
129            _ => {}
130        }
131    }
132    parts.push(StreamPart::Finish {
133        finish_reason: result.finish_reason.clone(),
134        usage: result.usage.clone(),
135        provider_metadata: result.provider_metadata.clone(),
136    });
137    parts
138}
139
140impl LanguageModelMiddleware for SimulateStreaming {
141    fn wrap_stream<'a>(
142        &'a self,
143        options: CallOptions,
144        _next: StreamNext<'a>,
145        ctx: MiddlewareContext<'a>,
146    ) -> BoxFuture<'a, Result<StreamResult, ProviderError>> {
147        Box::pin(async move {
148            let result = ctx.model.do_generate(options).await?;
149            let parts = simulate_parts(&result);
150            Ok(StreamResult {
151                stream: Box::pin(stream::iter(parts)),
152                request: result.request,
153                response: result.response,
154            })
155        })
156    }
157}