use crate::trace::Trace;
use crate::types::{SpanData, SpanType, SpanStatus};
pub struct Span<'a> {
trace: &'a Trace,
span_data: SpanData,
}
impl<'a> Span<'a> {
pub fn new(trace: &'a Trace, span_type: SpanType, name: impl Into<String>) -> Self {
Self {
trace,
span_data: SpanData {
id: uuid::Uuid::new_v4().to_string().replace("-", ""),
span_type,
name: name.into(),
input_tokens: 0,
output_tokens: 0,
cost: 0.0,
latency_ms: 0,
status: SpanStatus::Ok,
started_at: now_millis(),
ended_at: None,
model: None,
provider: None,
parent_span_id: None,
time_to_first_token_ms: None,
error_message: None,
input: None,
output: None,
metadata: None,
},
}
}
pub fn model(mut self, model: impl Into<String>) -> Self {
self.span_data.model = Some(model.into());
self
}
pub fn provider(mut self, provider: impl Into<String>) -> Self {
self.span_data.provider = Some(provider.into());
self
}
pub fn input_text(mut self, input: impl Into<String>) -> Self {
self.span_data.input = Some(input.into());
self
}
pub fn parent(mut self, parent_id: impl Into<String>) -> Self {
self.span_data.parent_span_id = Some(parent_id.into());
self
}
pub fn set_tokens(&mut self, input: i64, output: i64) {
self.span_data.input_tokens = input;
self.span_data.output_tokens = output;
}
pub fn set_cost(&mut self, cost: f64) {
self.span_data.cost = cost;
}
pub fn set_latency(&mut self, ms: i64) {
self.span_data.latency_ms = ms;
}
pub fn set_output(&mut self, output: impl Into<String>) {
self.span_data.output = Some(output.into());
}
pub fn set_error(&mut self, message: impl Into<String>) {
self.span_data.status = SpanStatus::Error;
self.span_data.error_message = Some(message.into());
}
pub fn set_time_to_first_token(&mut self, ms: i64) {
self.span_data.time_to_first_token_ms = Some(ms);
}
pub fn end(mut self) {
let end_time = now_millis();
self.span_data.ended_at = Some(end_time);
if self.span_data.latency_ms == 0 {
self.span_data.latency_ms = end_time - self.span_data.started_at;
}
self.trace.add_span(self.span_data);
}
pub fn data(&self) -> &SpanData {
&self.span_data
}
}
fn now_millis() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_millis() as i64
}