1use crate::trace::Trace;
2use crate::types::{SpanData, SpanType, SpanStatus};
3
4pub struct Span<'a> {
9 trace: &'a Trace,
10 span_data: SpanData,
11}
12
13impl<'a> Span<'a> {
14 pub fn new(trace: &'a Trace, span_type: SpanType, name: impl Into<String>) -> Self {
16 Self {
17 trace,
18 span_data: SpanData {
19 id: uuid::Uuid::new_v4().to_string().replace("-", ""),
20 span_type,
21 name: name.into(),
22 input_tokens: 0,
23 output_tokens: 0,
24 cost: 0.0,
25 latency_ms: 0,
26 status: SpanStatus::Ok,
27 started_at: now_millis(),
28 ended_at: None,
29 model: None,
30 provider: None,
31 parent_span_id: None,
32 time_to_first_token_ms: None,
33 error_message: None,
34 input: None,
35 output: None,
36 metadata: None,
37 },
38 }
39 }
40
41 pub fn model(mut self, model: impl Into<String>) -> Self {
44 self.span_data.model = Some(model.into());
45 self
46 }
47
48 pub fn provider(mut self, provider: impl Into<String>) -> Self {
49 self.span_data.provider = Some(provider.into());
50 self
51 }
52
53 pub fn input_text(mut self, input: impl Into<String>) -> Self {
54 self.span_data.input = Some(input.into());
55 self
56 }
57
58 pub fn parent(mut self, parent_id: impl Into<String>) -> Self {
59 self.span_data.parent_span_id = Some(parent_id.into());
60 self
61 }
62
63 pub fn set_tokens(&mut self, input: i64, output: i64) {
66 self.span_data.input_tokens = input;
67 self.span_data.output_tokens = output;
68 }
69
70 pub fn set_cost(&mut self, cost: f64) {
71 self.span_data.cost = cost;
72 }
73
74 pub fn set_latency(&mut self, ms: i64) {
75 self.span_data.latency_ms = ms;
76 }
77
78 pub fn set_output(&mut self, output: impl Into<String>) {
79 self.span_data.output = Some(output.into());
80 }
81
82 pub fn set_error(&mut self, message: impl Into<String>) {
83 self.span_data.status = SpanStatus::Error;
84 self.span_data.error_message = Some(message.into());
85 }
86
87 pub fn set_time_to_first_token(&mut self, ms: i64) {
88 self.span_data.time_to_first_token_ms = Some(ms);
89 }
90
91 pub fn end(mut self) {
94 let end_time = now_millis();
95 self.span_data.ended_at = Some(end_time);
96 if self.span_data.latency_ms == 0 {
97 self.span_data.latency_ms = end_time - self.span_data.started_at;
98 }
99 self.trace.add_span(self.span_data);
100 }
101
102 pub fn data(&self) -> &SpanData {
106 &self.span_data
107 }
108}
109
110fn now_millis() -> i64 {
111 std::time::SystemTime::now()
112 .duration_since(std::time::UNIX_EPOCH)
113 .unwrap()
114 .as_millis() as i64
115}