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    /// Execution resumed from a checkpoint
56    Resumed { step: usize, pending_nodes: Vec<String> },
57
58    /// Graph execution completed
59    Done { state: State, total_steps: usize },
60
61    /// Error occurred
62    Error { message: String, node: Option<String> },
63
64    /// Emitted when dynamic route dispatch occurs.
65    RouteDispatched { source: String, targets: Vec<String> },
66}
67
68impl StreamEvent {
69    /// Create a state event
70    pub fn state(state: State, step: usize) -> Self {
71        Self::State { state, step }
72    }
73
74    /// Create an updates event
75    pub fn updates(node: &str, updates: HashMap<String, Value>) -> Self {
76        Self::Updates { node: node.to_string(), updates }
77    }
78
79    /// Create a message event
80    pub fn message(node: &str, content: &str, is_final: bool) -> Self {
81        Self::Message { node: node.to_string(), content: content.to_string(), is_final }
82    }
83
84    /// Create a custom event
85    pub fn custom(node: &str, event_type: &str, data: Value) -> Self {
86        Self::Custom { node: node.to_string(), event_type: event_type.to_string(), data }
87    }
88
89    /// Create a debug event
90    pub fn debug(event_type: &str, data: Value) -> Self {
91        Self::Debug { event_type: event_type.to_string(), data }
92    }
93
94    /// Create a node start event
95    pub fn node_start(node: &str, step: usize) -> Self {
96        Self::NodeStart { node: node.to_string(), step }
97    }
98
99    /// Create a node end event
100    pub fn node_end(node: &str, step: usize, duration_ms: u64) -> Self {
101        Self::NodeEnd { node: node.to_string(), step, duration_ms }
102    }
103
104    /// Create a step complete event
105    pub fn step_complete(step: usize, nodes_executed: Vec<String>) -> Self {
106        Self::StepComplete { step, nodes_executed }
107    }
108
109    /// Create an interrupted event
110    pub fn interrupted(node: &str, message: &str) -> Self {
111        Self::Interrupted { node: node.to_string(), message: message.to_string() }
112    }
113
114    /// Create a resumed event
115    pub fn resumed(step: usize, pending_nodes: Vec<String>) -> Self {
116        Self::Resumed { step, pending_nodes }
117    }
118
119    /// Create a done event
120    pub fn done(state: State, total_steps: usize) -> Self {
121        Self::Done { state, total_steps }
122    }
123
124    /// Create an error event
125    pub fn error(message: &str, node: Option<&str>) -> Self {
126        Self::Error { message: message.to_string(), node: node.map(|s| s.to_string()) }
127    }
128
129    /// Create a route dispatched event
130    pub fn route_dispatched(source: &str, targets: Vec<String>) -> Self {
131        Self::RouteDispatched { source: source.to_string(), targets }
132    }
133}
adk_graph/stream.rs

adk_graph/
stream.rs
