Skip to main content

starweaver_model/
stream.rs

1//! Canonical model stream events.
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5
6use crate::message::ModelResponse;
7
8/// Stream event emitted by model adapters.
9#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
10#[serde(tag = "kind", rename_all = "snake_case")]
11pub enum ModelResponseStreamEvent {
12    /// A response part started.
13    PartStart(PartStart),
14    /// A response part delta arrived.
15    PartDelta(PartDelta),
16    /// A response part ended.
17    PartEnd(PartEnd),
18    /// Final response is available.
19    FinalResult(Box<ModelResponse>),
20}
21
22/// Part start event.
23#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
24pub struct PartStart {
25    /// Part index in response.
26    pub index: usize,
27    /// Part kind.
28    pub part_kind: String,
29}
30
31/// Part delta event.
32#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
33pub struct PartDelta {
34    /// Part index in response.
35    pub index: usize,
36    /// Typed delta payload.
37    #[serde(flatten)]
38    pub delta: StreamDelta,
39}
40
41impl PartDelta {
42    /// Build a text delta.
43    #[must_use]
44    pub fn text(index: usize, text: impl Into<String>) -> Self {
45        Self {
46            index,
47            delta: StreamDelta::Text { text: text.into() },
48        }
49    }
50
51    /// Build a thinking delta.
52    #[must_use]
53    pub fn thinking(index: usize, text: impl Into<String>) -> Self {
54        Self {
55            index,
56            delta: StreamDelta::Thinking { text: text.into() },
57        }
58    }
59
60    /// Return a text-only display projection.
61    #[must_use]
62    pub fn as_text(&self) -> String {
63        match &self.delta {
64            StreamDelta::Text { text }
65            | StreamDelta::Thinking { text }
66            | StreamDelta::ToolCallName { name: text }
67            | StreamDelta::ToolCallArguments {
68                arguments_delta: text,
69            } => text.clone(),
70            StreamDelta::NativePayload { payload } | StreamDelta::FileMetadata { payload } => {
71                payload.to_string()
72            }
73        }
74    }
75}
76
77/// Typed model stream delta payload.
78#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
79#[serde(tag = "delta_kind", rename_all = "snake_case")]
80pub enum StreamDelta {
81    /// Text output delta.
82    Text {
83        /// Text fragment.
84        text: String,
85    },
86    /// Thinking output delta.
87    Thinking {
88        /// Thinking fragment.
89        text: String,
90    },
91    /// Tool call name delta.
92    ToolCallName {
93        /// Tool name fragment or final value.
94        name: String,
95    },
96    /// Tool call argument delta.
97    ToolCallArguments {
98        /// Argument fragment.
99        arguments_delta: String,
100    },
101    /// Provider-native payload delta.
102    NativePayload {
103        /// Native payload fragment.
104        payload: Value,
105    },
106    /// File metadata delta.
107    FileMetadata {
108        /// File metadata fragment.
109        payload: Value,
110    },
111}
112
113/// Part end event.
114#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
115pub struct PartEnd {
116    /// Part index in response.
117    pub index: usize,
118    /// Part kind when the provider includes it or the parser can infer it.
119    #[serde(default, skip_serializing_if = "Option::is_none")]
120    pub part_kind: Option<String>,
121}
122
123impl PartEnd {
124    /// Build a part end event without a known part kind.
125    #[must_use]
126    pub const fn new(index: usize) -> Self {
127        Self {
128            index,
129            part_kind: None,
130        }
131    }
132
133    /// Build a part end event with an explicit part kind.
134    #[must_use]
135    pub fn with_kind(index: usize, part_kind: impl Into<String>) -> Self {
136        Self {
137            index,
138            part_kind: Some(part_kind.into()),
139        }
140    }
141}
142
143/// Lifecycle state of a streamed model response.
144#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
145#[serde(rename_all = "snake_case")]
146pub enum StreamLifecycle {
147    /// Events are still arriving.
148    #[default]
149    Incomplete,
150    /// Final response has been assembled.
151    Complete,
152    /// Stream ended by explicit cancellation or transport interruption.
153    Interrupted,
154}
155
156/// Lightweight stream assembly state for replay and lifecycle assertions.
157#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
158pub struct ModelStreamState {
159    /// Current lifecycle.
160    pub lifecycle: StreamLifecycle,
161    /// Number of started parts.
162    pub started_parts: usize,
163    /// Number of ended parts.
164    pub ended_parts: usize,
165    /// Final response when available.
166    #[serde(default, skip_serializing_if = "Option::is_none")]
167    pub final_response: Option<Box<ModelResponse>>,
168}
169
170impl ModelStreamState {
171    /// Apply one stream event.
172    pub fn apply(&mut self, event: &ModelResponseStreamEvent) {
173        match event {
174            ModelResponseStreamEvent::PartStart(_) => {
175                self.started_parts += 1;
176            }
177            ModelResponseStreamEvent::PartDelta(_) => {}
178            ModelResponseStreamEvent::PartEnd(_) => {
179                self.ended_parts += 1;
180            }
181            ModelResponseStreamEvent::FinalResult(response) => {
182                self.lifecycle = StreamLifecycle::Complete;
183                self.final_response = Some(response.clone());
184            }
185        }
186    }
187
188    /// Mark stream as interrupted.
189    pub fn interrupt(&mut self) {
190        self.lifecycle = StreamLifecycle::Interrupted;
191    }
192}