1use std::sync::Arc;
6
7pub use comindly_api::events::{Event, Kind, Level};
8
9pub trait EventSink: Send + Sync {
11 fn emit(&self, event: Event);
12}
13
14pub struct Discard;
16
17impl EventSink for Discard {
18 fn emit(&self, _: Event) {}
19}
20
21#[derive(Clone)]
23pub struct Reporter {
24 sink: Arc<dyn EventSink>,
25 step: Arc<std::sync::Mutex<Option<String>>>,
26}
27
28impl Reporter {
29 pub fn new(sink: Arc<dyn EventSink>) -> Self {
30 Reporter {
31 sink,
32 step: Default::default(),
33 }
34 }
35
36 fn send(&self, mut e: Event) {
37 e.step = self.step.lock().ok().and_then(|s| s.clone());
38 self.sink.emit(e);
39 }
40
41 pub fn step(&self, name: impl Into<String>) {
42 let name = name.into();
43 if let Ok(mut s) = self.step.lock() {
44 *s = Some(name.clone());
45 }
46 self.send(Event::new(Level::Info, Kind::Step, name));
47 }
48
49 pub fn done(&self, message: impl Into<String>) {
50 self.send(Event::new(Level::Info, Kind::Done, message));
51 }
52
53 pub fn note(&self, message: impl Into<String>) {
54 self.send(Event::new(Level::Info, Kind::Note, message));
55 }
56
57 pub fn warn(&self, message: impl Into<String>) {
58 self.send(Event::new(Level::Warn, Kind::Warning, message));
59 }
60
61 pub fn error(&self, message: impl Into<String>) {
62 self.send(Event::new(Level::Error, Kind::Error, message));
63 }
64
65 pub fn progress(&self, message: &str, done: u64, total: u64, usage: Option<Usage>) {
67 let mut e = Event::new(Level::Info, Kind::Progress, message);
68 e.done = Some(done);
69 e.total = Some(total);
70 if let Some(u) = usage {
71 e.tokens_in = Some(u.tokens_in);
72 e.tokens_out = Some(u.tokens_out);
73 e.usd = u.usd;
74 }
75 self.send(e);
76 }
77}
78
79#[derive(Debug, Clone, Copy, Default)]
81pub struct Usage {
82 pub tokens_in: u64,
83 pub tokens_out: u64,
84 pub usd: Option<f64>,
85}
86
87impl Usage {
88 pub fn priced(tokens_in: u64, tokens_out: u64, prices: Option<(f64, f64)>) -> Self {
89 Usage {
90 tokens_in,
91 tokens_out,
92 usd: prices.map(|(pi, po)| (tokens_in as f64 * pi + tokens_out as f64 * po) / 1e6),
93 }
94 }
95}
96
97impl std::fmt::Display for Usage {
98 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99 let k = |n: u64| {
100 if n >= 1000 {
101 format!("{:.1}k", n as f64 / 1000.0)
102 } else {
103 n.to_string()
104 }
105 };
106 write!(
107 f,
108 "{} in / {} out tokens",
109 k(self.tokens_in),
110 k(self.tokens_out)
111 )?;
112 match self.usd {
113 Some(u) if u < 0.01 => write!(f, " ≈ < $0.01"),
114 Some(u) => write!(f, " ≈ ${u:.2}"),
115 None => Ok(()),
116 }
117 }
118}