Skip to main content

adk_graph/
stream.rs

1//! Streaming types for graph execution
2
3use crate::state::State;
4use serde::Serialize;
5use serde_json::Value;
6use std::collections::HashMap;
7
8/// Stream mode options
9#[derive(Clone, Copy, Debug, Default)]
10pub enum StreamMode {
11    /// Full state after each super-step
12    #[default]
13    Values,
14    /// Only state changes
15    Updates,
16    /// LLM tokens and messages
17    Messages,
18    /// Custom events from nodes
19    Custom,
20    /// Debug information
21    Debug,
22}
23
24/// Events emitted during streaming
25#[derive(Clone, Debug, Serialize)]
26#[serde(tag = "type", rename_all = "snake_case")]
27pub enum StreamEvent {
28    /// State snapshot
29    State { state: State, step: usize },
30
31    /// State updates from a node
32    Updates { node: String, updates: HashMap<String, Value> },
33
34    /// Message/token from LLM
35    Message { node: String, content: String, is_final: bool },
36
37    /// Custom event from node
38    Custom { node: String, event_type: String, data: Value },
39
40    /// Debug event
41    Debug { event_type: String, data: Value },
42
43    /// Node started execution
44    NodeStart { node: String, step: usize },
45
46    /// Node completed execution
47    NodeEnd { node: String, step: usize, duration_ms: u64 },
48
49    /// Super-step completed
50    StepComplete { step: usize, nodes_executed: Vec<String> },
51
52    /// Execution was interrupted
53    Interrupted { node: String, message: String },
54
55    /// A node asked to pause, reported by `Node::execute_stream`.
56    ///
57    /// The streamed path yields events rather than a `NodeOutput`, so a node's own
58    /// interrupt request needs a way through. The executor turns this into the
59    /// pause and does not forward it, so a caller sees only `Interrupted`.
60    NodeInterrupt { node: String, message: String, data: Option<serde_json::Value> },
61
62    /// Execution resumed from a checkpoint
63    Resumed { step: usize, pending_nodes: Vec<String> },
64
65    /// Graph execution completed
66    Done { state: State, total_steps: usize },
67
68    /// Error occurred
69    Error { message: String, node: Option<String> },
70
71    /// Emitted when dynamic route dispatch occurs.
72    RouteDispatched { source: String, targets: Vec<String> },
73}
74
75impl StreamEvent {
76    /// Create a state event
77    pub fn state(state: State, step: usize) -> Self {
78        Self::State { state, step }
79    }
80
81    /// Create an updates event
82    pub fn updates(node: &str, updates: HashMap<String, Value>) -> Self {
83        Self::Updates { node: node.to_string(), updates }
84    }
85
86    /// Create a message event
87    pub fn message(node: &str, content: &str, is_final: bool) -> Self {
88        Self::Message { node: node.to_string(), content: content.to_string(), is_final }
89    }
90
91    /// Create a custom event
92    pub fn custom(node: &str, event_type: &str, data: Value) -> Self {
93        Self::Custom { node: node.to_string(), event_type: event_type.to_string(), data }
94    }
95
96    /// Create a debug event
97    pub fn debug(event_type: &str, data: Value) -> Self {
98        Self::Debug { event_type: event_type.to_string(), data }
99    }
100
101    /// Create a node start event
102    pub fn node_start(node: &str, step: usize) -> Self {
103        Self::NodeStart { node: node.to_string(), step }
104    }
105
106    /// Create a node end event
107    pub fn node_end(node: &str, step: usize, duration_ms: u64) -> Self {
108        Self::NodeEnd { node: node.to_string(), step, duration_ms }
109    }
110
111    /// Create a step complete event
112    pub fn step_complete(step: usize, nodes_executed: Vec<String>) -> Self {
113        Self::StepComplete { step, nodes_executed }
114    }
115
116    /// Create an interrupted event
117    pub fn interrupted(node: &str, message: &str) -> Self {
118        Self::Interrupted { node: node.to_string(), message: message.to_string() }
119    }
120
121    /// Create a resumed event
122    pub fn resumed(step: usize, pending_nodes: Vec<String>) -> Self {
123        Self::Resumed { step, pending_nodes }
124    }
125
126    /// Create a done event
127    pub fn done(state: State, total_steps: usize) -> Self {
128        Self::Done { state, total_steps }
129    }
130
131    /// Create an error event
132    pub fn error(message: &str, node: Option<&str>) -> Self {
133        Self::Error { message: message.to_string(), node: node.map(|s| s.to_string()) }
134    }
135
136    /// Create an event reporting that a node asked to pause.
137    pub fn node_interrupt(node: &str, message: &str, data: Option<serde_json::Value>) -> Self {
138        Self::NodeInterrupt { node: node.to_string(), message: message.to_string(), data }
139    }
140
141    /// Create a route dispatched event
142    pub fn route_dispatched(source: &str, targets: Vec<String>) -> Self {
143        Self::RouteDispatched { source: source.to_string(), targets }
144    }
145}