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

adk_graph/
stream.rs
