halo-agent 0.1.0

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

//! Settled model-call accounting for one Agent run. The history lifecycle uses
//! the total when writing the Chat projection; persistence stays outside Runtime.
use std::sync::{Arc, Mutex};

use crate::model::token::{self, Usage};

#[derive(Clone, Default)]
pub(crate) struct Ledger {
    calls: Arc<Mutex<Vec<Option<Usage>>>>,
    children: Arc<Mutex<Vec<Option<Usage>>>>,
}

impl Ledger {
    pub(crate) fn settle(&self, usage: Option<Usage>) {
        self.calls
            .lock()
            .expect("run usage lock poisoned")
            .push(usage);
    }

    pub(crate) fn settle_child(&self, usage: Option<Usage>) {
        if usage.is_some() {
            self.children
                .lock()
                .expect("child usage lock poisoned")
                .push(usage);
        }
    }

    pub(crate) fn refine_last(&self, usage: Option<Usage>) {
        if let Some(last) = self
            .calls
            .lock()
            .expect("run usage lock poisoned")
            .last_mut()
        {
            *last = usage;
        }
    }

    /// A restored ledger already contains all settled calls (or its parent owns
    /// their total). Add only newly known consumption, never replace that total
    /// with the single call whose after-processing is being resumed.
    pub(crate) fn refine_saved(&self, previous: Option<&Usage>, updated: &Usage) {
        fn added(previous: Option<u64>, updated: Option<u64>) -> Option<u64> {
            updated.map(|value| value.saturating_sub(previous.unwrap_or(0)))
        }
        self.settle(Some(Usage {
            input: added(previous.and_then(|v| v.input), updated.input),
            output: added(previous.and_then(|v| v.output), updated.output),
            total: added(previous.and_then(|v| v.total), updated.total),
            input_details: updated
                .input_details
                .as_ref()
                .map(|details| token::InputDetails {
                    cached: added(
                        previous.and_then(|v| v.input_details.as_ref()?.cached),
                        details.cached,
                    ),
                }),
            output_details: updated
                .output_details
                .as_ref()
                .map(|details| token::OutputDetails {
                    think: added(
                        previous.and_then(|v| v.output_details.as_ref()?.think),
                        details.think,
                    ),
                }),
            ..Default::default()
        }));
    }

    pub(crate) fn total(&self) -> Option<Usage> {
        let calls = self
            .calls
            .lock()
            .expect("run usage lock poisoned")
            .iter()
            .cloned()
            .fold(None, token::merge);
        self.children
            .lock()
            .expect("child usage lock poisoned")
            .iter()
            .cloned()
            .fold(calls, token::merge)
    }
}