1use std::sync::Arc;
6
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
10#[serde(rename_all = "lowercase")]
11pub enum Level {
12 Info,
13 Warn,
14 Error,
15}
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
18#[serde(rename_all = "lowercase")]
19pub enum Kind {
20 Step,
22 Done,
24 Note,
26 Warning,
27 Error,
28 Progress,
30}
31
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct Event {
34 pub timestamp: String,
35 pub level: Level,
36 pub kind: Kind,
37 pub message: String,
38 #[serde(skip_serializing_if = "Option::is_none")]
39 pub step: Option<String>,
40 #[serde(skip_serializing_if = "Option::is_none")]
42 pub job_id: Option<String>,
43 #[serde(skip_serializing_if = "Option::is_none")]
44 pub done: Option<u64>,
45 #[serde(skip_serializing_if = "Option::is_none")]
46 pub total: Option<u64>,
47 #[serde(skip_serializing_if = "Option::is_none")]
48 pub tokens_in: Option<u64>,
49 #[serde(skip_serializing_if = "Option::is_none")]
50 pub tokens_out: Option<u64>,
51 #[serde(skip_serializing_if = "Option::is_none")]
52 pub usd: Option<f64>,
53}
54
55impl Event {
56 pub fn new(level: Level, kind: Kind, message: impl Into<String>) -> Self {
57 Event {
58 timestamp: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
59 level,
60 kind,
61 message: message.into(),
62 step: None,
63 job_id: None,
64 done: None,
65 total: None,
66 tokens_in: None,
67 tokens_out: None,
68 usd: None,
69 }
70 }
71}
72
73pub trait EventSink: Send + Sync {
75 fn emit(&self, event: Event);
76}
77
78pub struct Discard;
80
81impl EventSink for Discard {
82 fn emit(&self, _: Event) {}
83}
84
85#[derive(Clone)]
87pub struct Reporter {
88 sink: Arc<dyn EventSink>,
89 step: Arc<std::sync::Mutex<Option<String>>>,
90}
91
92impl Reporter {
93 pub fn new(sink: Arc<dyn EventSink>) -> Self {
94 Reporter {
95 sink,
96 step: Default::default(),
97 }
98 }
99
100 fn send(&self, mut e: Event) {
101 e.step = self.step.lock().ok().and_then(|s| s.clone());
102 self.sink.emit(e);
103 }
104
105 pub fn step(&self, name: impl Into<String>) {
106 let name = name.into();
107 if let Ok(mut s) = self.step.lock() {
108 *s = Some(name.clone());
109 }
110 self.send(Event::new(Level::Info, Kind::Step, name));
111 }
112
113 pub fn done(&self, message: impl Into<String>) {
114 self.send(Event::new(Level::Info, Kind::Done, message));
115 }
116
117 pub fn note(&self, message: impl Into<String>) {
118 self.send(Event::new(Level::Info, Kind::Note, message));
119 }
120
121 pub fn warn(&self, message: impl Into<String>) {
122 self.send(Event::new(Level::Warn, Kind::Warning, message));
123 }
124
125 pub fn error(&self, message: impl Into<String>) {
126 self.send(Event::new(Level::Error, Kind::Error, message));
127 }
128
129 pub fn progress(&self, message: &str, done: u64, total: u64, usage: Option<Usage>) {
131 let mut e = Event::new(Level::Info, Kind::Progress, message);
132 e.done = Some(done);
133 e.total = Some(total);
134 if let Some(u) = usage {
135 e.tokens_in = Some(u.tokens_in);
136 e.tokens_out = Some(u.tokens_out);
137 e.usd = u.usd;
138 }
139 self.send(e);
140 }
141}
142
143#[derive(Debug, Clone, Copy, Default)]
145pub struct Usage {
146 pub tokens_in: u64,
147 pub tokens_out: u64,
148 pub usd: Option<f64>,
149}
150
151impl Usage {
152 pub fn priced(tokens_in: u64, tokens_out: u64, prices: Option<(f64, f64)>) -> Self {
153 Usage {
154 tokens_in,
155 tokens_out,
156 usd: prices.map(|(pi, po)| (tokens_in as f64 * pi + tokens_out as f64 * po) / 1e6),
157 }
158 }
159}
160
161impl std::fmt::Display for Usage {
162 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
163 let k = |n: u64| {
164 if n >= 1000 {
165 format!("{:.1}k", n as f64 / 1000.0)
166 } else {
167 n.to_string()
168 }
169 };
170 write!(
171 f,
172 "{} in / {} out tokens",
173 k(self.tokens_in),
174 k(self.tokens_out)
175 )?;
176 match self.usd {
177 Some(u) if u < 0.01 => write!(f, " ≈ < $0.01"),
178 Some(u) => write!(f, " ≈ ${u:.2}"),
179 None => Ok(()),
180 }
181 }
182}