halo-agent 0.1.1

Halo Agent and Workflow execution kernel
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

//! Agent-owned final-result processing, separate from an internal Model Step.
use crate::runtime::agent::{
    hooks::{Hooks, Stage},
    model_steps,
    node::Step,
};
use crate::{Cancellation, Error, Message, event::Event};
use serde::{Deserialize, Serialize};

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct State {
    /// after, complete, or committed; no separate Hook execution is created.
    pub stage: String,
    /// Number of successfully completed middleware methods in this stage.
    pub idx: usize,
    /// Last successful transformation, never an unfinished Hook's working copy.
    pub message: Message,
    pub finish_reason: Option<String>,
}
impl State {
    pub(crate) fn new(message: Message, finish_reason: Option<String>) -> Self {
        Self {
            stage: "after".into(),
            idx: 0,
            message,
            finish_reason,
        }
    }
    pub(crate) fn validate(&self, count: usize) -> Result<(), Error> {
        if !matches!(self.stage.as_str(), "after" | "complete" | "committed")
            || self.idx > count
            || (self.stage != "after" && self.idx != count)
        {
            return Err(Error::new(
                "INVALID_STATE",
                "invalid Agent completion phase",
            ));
        }
        Ok(())
    }
    pub(crate) fn pending(&self) -> bool {
        self.stage != "committed"
    }
}

pub(super) async fn model(
    hooks: &Hooks,
    mut saved: Step,
    mut commit: impl FnMut(Step) -> Result<(), Error>,
    cancellation: &Cancellation,
) -> Result<Event, Error> {
    let message = saved
        .message
        .clone()
        .ok_or_else(|| Error::new("INVALID_STATE", "missing Model result"))?;
    let result = Box::pin(hooks.after_from(
        Stage::Model,
        message,
        saved.idx.unwrap_or(0),
        |idx, message| {
            saved.idx = Some(idx);
            saved.message = Some(message);
            commit(saved.clone())
        },
        cancellation,
    ))
    .await;
    match result {
        Ok(message) => {
            saved.message = Some(message);
            saved.stage = Some("committed".into());
            saved.status = "succeeded".into();
            saved.error = None;
            commit(saved.clone())?;
            model_steps::event(&saved)
        }
        Err(error) => {
            model_steps::stopped(std::slice::from_mut(&mut saved), &error);
            commit(saved)?;
            Err(error)
        }
    }
}

pub(super) async fn agent(
    hooks: &Hooks,
    mut saved: State,
    mut commit: impl FnMut(State) -> Result<(), Error>,
    cancellation: &Cancellation,
) -> Result<Event, Error> {
    if saved.stage == "after" {
        let message = Box::pin(hooks.after_from(
            Stage::Execute,
            saved.message.clone(),
            saved.idx,
            |idx, message| {
                saved.idx = idx;
                saved.message = message;
                commit(saved.clone())
            },
            cancellation,
        ))
        .await?;
        saved.message = message;
        saved.stage = "complete".into();
        commit(saved.clone())?;
    }
    Ok(Event::Complete {
        usage: saved.message.usage.clone(),
        message: saved.message,
        finish_reason: saved.finish_reason,
    })
}

pub(super) fn has_calls(message: &Message) -> bool {
    message
        .content
        .iter()
        .any(|part| part.r#type == "tool_call")
}

/// Cancellation observes the same acknowledged result as restoration, not an
/// unfinished middleware working copy or the stale pre-transform raw output.
pub(super) fn partial(completion: Option<&State>, steps: &[Step]) -> Option<Message> {
    completion
        .filter(|state| state.pending())
        .map(|state| state.message.clone())
        .or_else(|| model_steps::pending(steps).and_then(|step| step.message.clone()))
}