bloop_sdk/
trace.rs

1use std::sync::Mutex;
2use crate::types::{TraceData, TraceStatus, SpanData};
3
4/// A live LLM trace that collects spans.
5///
6/// Traces use interior mutability via Mutex so that spans can add themselves
7/// to the trace without requiring exclusive (&mut) access. This enables the
8/// ergonomic pattern:
9///
10/// ```rust,ignore
11/// let trace = client.start_trace("chat");
12/// let mut span = Span::new(&trace, SpanType::Generation, "llm");
13/// span.end();  // adds to trace internally
14/// trace.end();
15/// ```
16pub struct Trace {
17    id: String,
18    name: String,
19    status: Mutex<TraceStatus>,
20    started_at: i64,
21    ended_at: Mutex<Option<i64>>,
22    session_id: Option<String>,
23    user_id: Option<String>,
24    input: Option<String>,
25    output: Mutex<Option<String>>,
26    metadata: Option<serde_json::Value>,
27    prompt_name: Option<String>,
28    prompt_version: Option<String>,
29    spans: Mutex<Vec<SpanData>>,
30}
31
32impl Trace {
33    /// Create a new trace with the given name.
34    pub fn new(name: impl Into<String>) -> Self {
35        Self {
36            id: uuid::Uuid::new_v4().to_string().replace("-", ""),
37            name: name.into(),
38            status: Mutex::new(TraceStatus::Running),
39            started_at: now_millis(),
40            ended_at: Mutex::new(None),
41            session_id: None,
42            user_id: None,
43            input: None,
44            output: Mutex::new(None),
45            metadata: None,
46            prompt_name: None,
47            prompt_version: None,
48            spans: Mutex::new(Vec::new()),
49        }
50    }
51
52    // ── Builder-pattern setters (consume self) ──
53
54    pub fn session_id(mut self, id: impl Into<String>) -> Self {
55        self.session_id = Some(id.into());
56        self
57    }
58
59    pub fn user_id(mut self, id: impl Into<String>) -> Self {
60        self.user_id = Some(id.into());
61        self
62    }
63
64    pub fn input_text(mut self, input: impl Into<String>) -> Self {
65        self.input = Some(input.into());
66        self
67    }
68
69    pub fn prompt_name(mut self, name: impl Into<String>) -> Self {
70        self.prompt_name = Some(name.into());
71        self
72    }
73
74    pub fn prompt_version(mut self, version: impl Into<String>) -> Self {
75        self.prompt_version = Some(version.into());
76        self
77    }
78
79    // ── Mutating methods ──
80
81    pub fn set_output(&self, output: impl Into<String>) {
82        *self.output.lock().unwrap() = Some(output.into());
83    }
84
85    pub fn set_status(&self, status: TraceStatus) {
86        *self.status.lock().unwrap() = status;
87    }
88
89    pub fn end(&self) {
90        let mut ended = self.ended_at.lock().unwrap();
91        if ended.is_none() {
92            *ended = Some(now_millis());
93        }
94        let mut status = self.status.lock().unwrap();
95        if *status == TraceStatus::Running {
96            *status = TraceStatus::Completed;
97        }
98    }
99
100    // ── Internal: called by Span::end() ──
101
102    pub(crate) fn add_span(&self, span_data: SpanData) {
103        self.spans.lock().unwrap().push(span_data);
104    }
105
106    // ── Accessors ──
107
108    pub fn id(&self) -> &str {
109        &self.id
110    }
111
112    pub fn name(&self) -> &str {
113        &self.name
114    }
115
116    pub fn started_at(&self) -> i64 {
117        self.started_at
118    }
119
120    pub fn span_count(&self) -> usize {
121        self.spans.lock().unwrap().len()
122    }
123
124    pub fn spans(&self) -> Vec<SpanData> {
125        self.spans.lock().unwrap().clone()
126    }
127
128    /// Convert to owned TraceData for serialization/sending.
129    pub fn to_data(&self) -> TraceData {
130        TraceData {
131            id: self.id.clone(),
132            name: self.name.clone(),
133            status: *self.status.lock().unwrap(),
134            started_at: self.started_at,
135            ended_at: *self.ended_at.lock().unwrap(),
136            spans: self.spans.lock().unwrap().clone(),
137            session_id: self.session_id.clone(),
138            user_id: self.user_id.clone(),
139            input: self.input.clone(),
140            output: self.output.lock().unwrap().clone(),
141            metadata: self.metadata.clone(),
142            prompt_name: self.prompt_name.clone(),
143            prompt_version: self.prompt_version.clone(),
144        }
145    }
146}
147
148impl std::fmt::Debug for Trace {
149    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
150        f.debug_struct("Trace")
151            .field("id", &self.id)
152            .field("name", &self.name)
153            .field("started_at", &self.started_at)
154            .field("span_count", &self.span_count())
155            .finish()
156    }
157}
158
159fn now_millis() -> i64 {
160    std::time::SystemTime::now()
161        .duration_since(std::time::UNIX_EPOCH)
162        .unwrap()
163        .as_millis() as i64
164}
bloop_sdk/trace.rs

bloop_sdk/
trace.rs
