use serde::{Deserialize, Serialize};
use crate::{
AgentResult,
executor::AgentExecutionNode,
stream::{AgentStreamEvent, AgentStreamRecord},
};
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum AgentIterationKind {
RunStart,
NodeStart,
NodeComplete,
Custom,
ModelRequest,
ModelStream,
ModelResponse,
Checkpoint,
Suspended,
ToolCall,
ToolReturn,
OutputRetry,
SteeringGuard,
RunComplete,
RunFailed,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct AgentIterationStep {
pub index: usize,
pub stream_sequence: usize,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub run_step: Option<usize>,
pub kind: AgentIterationKind,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub node: Option<AgentExecutionNode>,
}
impl AgentIterationStep {
const fn from_record(index: usize, record: &AgentStreamRecord) -> Self {
let (kind, run_step, node) = match &record.event {
AgentStreamEvent::RunStart { .. } => (AgentIterationKind::RunStart, None, None),
AgentStreamEvent::NodeStart { node, step, .. } => {
(AgentIterationKind::NodeStart, Some(*step), Some(*node))
}
AgentStreamEvent::NodeComplete { node, step, .. } => {
(AgentIterationKind::NodeComplete, Some(*step), Some(*node))
}
AgentStreamEvent::Custom { .. } => (AgentIterationKind::Custom, None, None),
AgentStreamEvent::ModelRequest { step } => {
(AgentIterationKind::ModelRequest, Some(*step), None)
}
AgentStreamEvent::ModelStream { step, .. } => {
(AgentIterationKind::ModelStream, Some(*step), None)
}
AgentStreamEvent::ModelResponse { step, .. } => {
(AgentIterationKind::ModelResponse, Some(*step), None)
}
AgentStreamEvent::Checkpoint { node, step } => {
(AgentIterationKind::Checkpoint, Some(*step), Some(*node))
}
AgentStreamEvent::Suspended { node, .. } => {
(AgentIterationKind::Suspended, None, Some(*node))
}
AgentStreamEvent::ToolCall { step, .. } => {
(AgentIterationKind::ToolCall, Some(*step), None)
}
AgentStreamEvent::ToolReturn { step, .. } => {
(AgentIterationKind::ToolReturn, Some(*step), None)
}
AgentStreamEvent::OutputRetry { .. } => (AgentIterationKind::OutputRetry, None, None),
AgentStreamEvent::SteeringGuard { step, .. } => {
(AgentIterationKind::SteeringGuard, Some(*step), None)
}
AgentStreamEvent::RunComplete { .. } => (AgentIterationKind::RunComplete, None, None),
AgentStreamEvent::RunFailed { .. } => (AgentIterationKind::RunFailed, None, None),
};
Self {
index,
stream_sequence: record.sequence,
run_step,
kind,
node,
}
}
}
#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
pub struct AgentIterationTrace {
pub steps: Vec<AgentIterationStep>,
}
impl AgentIterationTrace {
#[must_use]
pub fn from_stream_records(records: &[AgentStreamRecord]) -> Self {
Self {
steps: records
.iter()
.enumerate()
.map(|(index, record)| AgentIterationStep::from_record(index, record))
.collect(),
}
}
#[must_use]
pub fn steps(&self) -> &[AgentIterationStep] {
&self.steps
}
#[must_use]
pub fn is_complete(&self) -> bool {
self.steps
.iter()
.any(|step| step.kind == AgentIterationKind::RunComplete)
}
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct AgentIterResult {
pub result: AgentResult,
pub iterations: AgentIterationTrace,
pub events: Vec<AgentStreamRecord>,
}
impl AgentIterResult {
#[must_use]
pub const fn result(&self) -> &AgentResult {
&self.result
}
#[must_use]
pub const fn iterations(&self) -> &AgentIterationTrace {
&self.iterations
}
#[must_use]
pub fn events(&self) -> &[AgentStreamRecord] {
&self.events
}
}