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 RunFailed,
45}
46
47#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
49pub struct AgentIterationStep {
50 pub index: usize,
52 pub stream_sequence: usize,
54 #[serde(default, skip_serializing_if = "Option::is_none")]
56 pub run_step: Option<usize>,
57 pub kind: AgentIterationKind,
59 #[serde(default, skip_serializing_if = "Option::is_none")]
61 pub node: Option<AgentExecutionNode>,
62}
63
64impl AgentIterationStep {
65 const fn from_record(index: usize, record: &AgentStreamRecord) -> Self {
66 let (kind, run_step, node) = match &record.event {
67 AgentStreamEvent::RunStart { .. } => (AgentIterationKind::RunStart, None, None),
68 AgentStreamEvent::NodeStart { node, step, .. } => {
69 (AgentIterationKind::NodeStart, Some(*step), Some(*node))
70 }
71 AgentStreamEvent::NodeComplete { node, step, .. } => {
72 (AgentIterationKind::NodeComplete, Some(*step), Some(*node))
73 }
74 AgentStreamEvent::Custom { .. } => (AgentIterationKind::Custom, None, None),
75 AgentStreamEvent::ModelRequest { step } => {
76 (AgentIterationKind::ModelRequest, Some(*step), None)
77 }
78 AgentStreamEvent::ModelStream { step, .. } => {
79 (AgentIterationKind::ModelStream, Some(*step), None)
80 }
81 AgentStreamEvent::ModelResponse { step, .. } => {
82 (AgentIterationKind::ModelResponse, Some(*step), None)
83 }
84 AgentStreamEvent::Checkpoint { node, step } => {
85 (AgentIterationKind::Checkpoint, Some(*step), Some(*node))
86 }
87 AgentStreamEvent::Suspended { node, .. } => {
88 (AgentIterationKind::Suspended, None, Some(*node))
89 }
90 AgentStreamEvent::ToolCall { step, .. } => {
91 (AgentIterationKind::ToolCall, Some(*step), None)
92 }
93 AgentStreamEvent::ToolReturn { step, .. } => {
94 (AgentIterationKind::ToolReturn, Some(*step), None)
95 }
96 AgentStreamEvent::OutputRetry { .. } => (AgentIterationKind::OutputRetry, None, None),
97 AgentStreamEvent::SteeringGuard { step, .. } => {
98 (AgentIterationKind::SteeringGuard, Some(*step), None)
99 }
100 AgentStreamEvent::RunComplete { .. } => (AgentIterationKind::RunComplete, None, None),
101 AgentStreamEvent::RunFailed { .. } => (AgentIterationKind::RunFailed, None, None),
102 };
103 Self {
104 index,
105 stream_sequence: record.sequence,
106 run_step,
107 kind,
108 node,
109 }
110 }
111}
112
113#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
115pub struct AgentIterationTrace {
116 pub steps: Vec<AgentIterationStep>,
118}
119
120impl AgentIterationTrace {
121 #[must_use]
123 pub fn from_stream_records(records: &[AgentStreamRecord]) -> Self {
124 Self {
125 steps: records
126 .iter()
127 .enumerate()
128 .map(|(index, record)| AgentIterationStep::from_record(index, record))
129 .collect(),
130 }
131 }
132
133 #[must_use]
135 pub fn steps(&self) -> &[AgentIterationStep] {
136 &self.steps
137 }
138
139 #[must_use]
141 pub fn is_complete(&self) -> bool {
142 self.steps
143 .iter()
144 .any(|step| step.kind == AgentIterationKind::RunComplete)
145 }
146}
147
148#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
150pub struct AgentIterResult {
151 pub result: AgentResult,
153 pub iterations: AgentIterationTrace,
155 pub events: Vec<AgentStreamRecord>,
157}
158
159impl AgentIterResult {
160 #[must_use]
162 pub const fn result(&self) -> &AgentResult {
163 &self.result
164 }
165
166 #[must_use]
168 pub const fn iterations(&self) -> &AgentIterationTrace {
169 &self.iterations
170 }
171
172 #[must_use]
174 pub fn events(&self) -> &[AgentStreamRecord] {
175 &self.events
176 }
177}