1use serde::{Deserialize, Serialize};
4use serde_json::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 FinalResult(Box<ModelResponse>),
20}
21
22#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
24pub struct PartStart {
25 pub index: usize,
27 pub part_kind: String,
29}
30
31#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
33pub struct PartDelta {
34 pub index: usize,
36 #[serde(flatten)]
38 pub delta: StreamDelta,
39}
40
41impl PartDelta {
42 #[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 #[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 #[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#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
79#[serde(tag = "delta_kind", rename_all = "snake_case")]
80pub enum StreamDelta {
81 Text {
83 text: String,
85 },
86 Thinking {
88 text: String,
90 },
91 ToolCallName {
93 name: String,
95 },
96 ToolCallArguments {
98 arguments_delta: String,
100 },
101 NativePayload {
103 payload: Value,
105 },
106 FileMetadata {
108 payload: Value,
110 },
111}
112
113#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
115pub struct PartEnd {
116 pub index: usize,
118 #[serde(default, skip_serializing_if = "Option::is_none")]
120 pub part_kind: Option<String>,
121}
122
123impl PartEnd {
124 #[must_use]
126 pub const fn new(index: usize) -> Self {
127 Self {
128 index,
129 part_kind: None,
130 }
131 }
132
133 #[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#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
145#[serde(rename_all = "snake_case")]
146pub enum StreamLifecycle {
147 #[default]
149 Incomplete,
150 Complete,
152 Interrupted,
154}
155
156#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
158pub struct ModelStreamState {
159 pub lifecycle: StreamLifecycle,
161 pub started_parts: usize,
163 pub ended_parts: usize,
165 #[serde(default, skip_serializing_if = "Option::is_none")]
167 pub final_response: Option<Box<ModelResponse>>,
168}
169
170impl ModelStreamState {
171 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 pub fn interrupt(&mut self) {
190 self.lifecycle = StreamLifecycle::Interrupted;
191 }
192}