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
65impl StreamEvent {
66    /// Create a state event
67    pub fn state(state: State, step: usize) -> Self {
68        Self::State { state, step }
69    }
70
71    /// Create an updates event
72    pub fn updates(node: &str, updates: HashMap<String, Value>) -> Self {
73        Self::Updates { node: node.to_string(), updates }
74    }
75
76    /// Create a message event
77    pub fn message(node: &str, content: &str, is_final: bool) -> Self {
78        Self::Message { node: node.to_string(), content: content.to_string(), is_final }
79    }
80
81    /// Create a custom event
82    pub fn custom(node: &str, event_type: &str, data: Value) -> Self {
83        Self::Custom { node: node.to_string(), event_type: event_type.to_string(), data }
84    }
85
86    /// Create a debug event
87    pub fn debug(event_type: &str, data: Value) -> Self {
88        Self::Debug { event_type: event_type.to_string(), data }
89    }
90
91    /// Create a node start event
92    pub fn node_start(node: &str, step: usize) -> Self {
93        Self::NodeStart { node: node.to_string(), step }
94    }
95
96    /// Create a node end event
97    pub fn node_end(node: &str, step: usize, duration_ms: u64) -> Self {
98        Self::NodeEnd { node: node.to_string(), step, duration_ms }
99    }
100
101    /// Create a step complete event
102    pub fn step_complete(step: usize, nodes_executed: Vec<String>) -> Self {
103        Self::StepComplete { step, nodes_executed }
104    }
105
106    /// Create an interrupted event
107    pub fn interrupted(node: &str, message: &str) -> Self {
108        Self::Interrupted { node: node.to_string(), message: message.to_string() }
109    }
110
111    /// Create a resumed event
112    pub fn resumed(step: usize, pending_nodes: Vec<String>) -> Self {
113        Self::Resumed { step, pending_nodes }
114    }
115
116    /// Create a done event
117    pub fn done(state: State, total_steps: usize) -> Self {
118        Self::Done { state, total_steps }
119    }
120
121    /// Create an error event
122    pub fn error(message: &str, node: Option<&str>) -> Self {
123        Self::Error { message: message.to_string(), node: node.map(|s| s.to_string()) }
124    }
125}
adk_graph/stream.rs

adk_graph/
stream.rs
