1use serde::{Deserialize, Serialize};
4use serde_json::{Map, Value};
5
6use crate::message::ModelResponse;
7
8#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
10#[serde(tag = "kind", rename_all = "snake_case")]
11pub enum ModelResponseStreamEvent {
12 PartStart(PartStart),
14 PartDelta(PartDelta),
16 PartEnd(PartEnd),
18 Diagnostic(StreamDiagnostic),
20 FinalResult(Box<ModelResponse>),
22}
23
24#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
26pub struct StreamDiagnostic {
27 pub kind: String,
29 #[serde(default)]
31 pub payload: Value,
32 #[serde(default, skip_serializing_if = "Map::is_empty")]
34 pub metadata: Map<String, Value>,
35}
36
37impl StreamDiagnostic {
38 #[must_use]
40 pub fn new(kind: impl Into<String>, payload: Value) -> Self {
41 Self {
42 kind: kind.into(),
43 payload,
44 metadata: Map::new(),
45 }
46 }
47}
48
49#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
51pub struct PartStart {
52 pub index: usize,
54 pub part_kind: String,
56}
57
58#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
60pub struct PartDelta {
61 pub index: usize,
63 #[serde(flatten)]
65 pub delta: StreamDelta,
66}
67
68impl PartDelta {
69 #[must_use]
71 pub fn text(index: usize, text: impl Into<String>) -> Self {
72 Self {
73 index,
74 delta: StreamDelta::Text { text: text.into() },
75 }
76 }
77
78 #[must_use]
80 pub fn thinking(index: usize, text: impl Into<String>) -> Self {
81 Self {
82 index,
83 delta: StreamDelta::Thinking { text: text.into() },
84 }
85 }
86
87 #[must_use]
89 pub fn as_text(&self) -> String {
90 match &self.delta {
91 StreamDelta::Text { text }
92 | StreamDelta::Thinking { text }
93 | StreamDelta::ToolCallName { name: text }
94 | StreamDelta::ToolCallArguments {
95 arguments_delta: text,
96 } => text.clone(),
97 StreamDelta::NativePayload { payload } | StreamDelta::FileMetadata { payload } => {
98 payload.to_string()
99 }
100 }
101 }
102}
103
104#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
106#[serde(tag = "delta_kind", rename_all = "snake_case")]
107pub enum StreamDelta {
108 Text {
110 text: String,
112 },
113 Thinking {
115 text: String,
117 },
118 ToolCallName {
120 name: String,
122 },
123 ToolCallArguments {
125 arguments_delta: String,
127 },
128 NativePayload {
130 payload: Value,
132 },
133 FileMetadata {
135 payload: Value,
137 },
138}
139
140#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
142pub struct PartEnd {
143 pub index: usize,
145 #[serde(default, skip_serializing_if = "Option::is_none")]
147 pub part_kind: Option<String>,
148}
149
150impl PartEnd {
151 #[must_use]
153 pub const fn new(index: usize) -> Self {
154 Self {
155 index,
156 part_kind: None,
157 }
158 }
159
160 #[must_use]
162 pub fn with_kind(index: usize, part_kind: impl Into<String>) -> Self {
163 Self {
164 index,
165 part_kind: Some(part_kind.into()),
166 }
167 }
168}
169
170#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
172#[serde(rename_all = "snake_case")]
173pub enum StreamLifecycle {
174 #[default]
176 Incomplete,
177 Complete,
179 Interrupted,
181}
182
183#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
185pub struct ModelStreamState {
186 pub lifecycle: StreamLifecycle,
188 pub started_parts: usize,
190 pub ended_parts: usize,
192 #[serde(default, skip_serializing_if = "Option::is_none")]
194 pub final_response: Option<Box<ModelResponse>>,
195}
196
197impl ModelStreamState {
198 pub fn apply(&mut self, event: &ModelResponseStreamEvent) {
200 match event {
201 ModelResponseStreamEvent::PartStart(_) => {
202 self.started_parts += 1;
203 }
204 ModelResponseStreamEvent::PartDelta(_) | ModelResponseStreamEvent::Diagnostic(_) => {}
205 ModelResponseStreamEvent::PartEnd(_) => {
206 self.ended_parts += 1;
207 }
208 ModelResponseStreamEvent::FinalResult(response) => {
209 self.lifecycle = StreamLifecycle::Complete;
210 self.final_response = Some(response.clone());
211 }
212 }
213 }
214
215 pub const fn interrupt(&mut self) {
217 self.lifecycle = StreamLifecycle::Interrupted;
218 }
219}