1use serde::{Deserialize, Serialize};
4
5use crate::{
6 AgentResult,
7 executor::AgentExecutionNode,
8 stream::{AgentStreamEvent, AgentStreamRecord},
9};
10
11#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
13#[serde(rename_all = "snake_case")]
14pub enum AgentIterationKind {
15 RunStart,
17 NodeStart,
19 NodeComplete,
21 Custom,
23 ModelRequest,
25 ModelStream,
27 ModelResponse,
29 Checkpoint,
31 Suspended,
33 ToolCall,
35 ToolReturn,
37 OutputRetry,
39 SteeringGuard,
41 RunComplete,
43 RunCancelled,
45 RunFailed,
47}
48
49#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
51pub struct AgentIterationStep {
52 pub index: usize,
54 pub stream_sequence: usize,
56 #[serde(default, skip_serializing_if = "Option::is_none")]
58 pub run_step: Option<usize>,
59 pub kind: AgentIterationKind,
61 #[serde(default, skip_serializing_if = "Option::is_none")]
63 pub node: Option<AgentExecutionNode>,
64}
65
66impl AgentIterationStep {
67 const fn from_record(index: usize, record: &AgentStreamRecord) -> Self {
68 let (kind, run_step, node) = match &record.event {
69 AgentStreamEvent::RunStart { .. } => (AgentIterationKind::RunStart, None, None),
70 AgentStreamEvent::NodeStart { node, step, .. } => {
71 (AgentIterationKind::NodeStart, Some(*step), Some(*node))
72 }
73 AgentStreamEvent::NodeComplete { node, step, .. } => {
74 (AgentIterationKind::NodeComplete, Some(*step), Some(*node))
75 }
76 AgentStreamEvent::Custom { .. } => (AgentIterationKind::Custom, None, None),
77 AgentStreamEvent::ModelRequest { step } => {
78 (AgentIterationKind::ModelRequest, Some(*step), None)
79 }
80 AgentStreamEvent::ModelStream { step, .. } => {
81 (AgentIterationKind::ModelStream, Some(*step), None)
82 }
83 AgentStreamEvent::ModelResponse { step, .. } => {
84 (AgentIterationKind::ModelResponse, Some(*step), None)
85 }
86 AgentStreamEvent::Checkpoint { node, step } => {
87 (AgentIterationKind::Checkpoint, Some(*step), Some(*node))
88 }
89 AgentStreamEvent::Suspended { node, .. } => {
90 (AgentIterationKind::Suspended, None, Some(*node))
91 }
92 AgentStreamEvent::ToolCall { step, .. } => {
93 (AgentIterationKind::ToolCall, Some(*step), None)
94 }
95 AgentStreamEvent::ToolReturn { step, .. } => {
96 (AgentIterationKind::ToolReturn, Some(*step), None)
97 }
98 AgentStreamEvent::OutputRetry { .. } => (AgentIterationKind::OutputRetry, None, None),
99 AgentStreamEvent::SteeringGuard { step, .. } => {
100 (AgentIterationKind::SteeringGuard, Some(*step), None)
101 }
102 AgentStreamEvent::RunComplete { .. } => (AgentIterationKind::RunComplete, None, None),
103 AgentStreamEvent::RunCancelled { .. } => (AgentIterationKind::RunCancelled, None, None),
104 AgentStreamEvent::RunFailed { .. } => (AgentIterationKind::RunFailed, None, None),
105 };
106 Self {
107 index,
108 stream_sequence: record.sequence,
109 run_step,
110 kind,
111 node,
112 }
113 }
114}
115
116#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
118pub struct AgentIterationTrace {
119 pub steps: Vec<AgentIterationStep>,
121}
122
123impl AgentIterationTrace {
124 #[must_use]
126 pub fn from_stream_records(records: &[AgentStreamRecord]) -> Self {
127 Self {
128 steps: records
129 .iter()
130 .enumerate()
131 .map(|(index, record)| AgentIterationStep::from_record(index, record))
132 .collect(),
133 }
134 }
135
136 #[must_use]
138 pub fn steps(&self) -> &[AgentIterationStep] {
139 &self.steps
140 }
141
142 #[must_use]
144 pub fn is_complete(&self) -> bool {
145 self.steps
146 .iter()
147 .any(|step| step.kind == AgentIterationKind::RunComplete)
148 }
149}
150
151#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
153pub struct AgentIterResult {
154 pub result: AgentResult,
156 pub iterations: AgentIterationTrace,
158 pub events: Vec<AgentStreamRecord>,
160}
161
162impl AgentIterResult {
163 #[must_use]
165 pub const fn result(&self) -> &AgentResult {
166 &self.result
167 }
168
169 #[must_use]
171 pub const fn iterations(&self) -> &AgentIterationTrace {
172 &self.iterations
173 }
174
175 #[must_use]
177 pub fn events(&self) -> &[AgentStreamRecord] {
178 &self.events
179 }
180}