1use crate::types::{Cost, Latency, Metadata, Provider, TokenUsage, TraceId, SpanId};
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct LlmSpan {
14 pub span_id: SpanId,
16 pub trace_id: TraceId,
18 pub parent_span_id: Option<SpanId>,
20 pub name: String,
22 pub provider: Provider,
24 pub model: String,
26 pub input: LlmInput,
28 pub output: Option<LlmOutput>,
30 pub token_usage: Option<TokenUsage>,
32 pub cost: Option<Cost>,
34 pub latency: Latency,
36 pub metadata: Metadata,
38 pub status: SpanStatus,
40 #[serde(default)]
42 pub attributes: HashMap<String, serde_json::Value>,
43 #[serde(default)]
45 pub events: Vec<SpanEvent>,
46}
47
48#[derive(Debug, Clone, Serialize, Deserialize)]
50#[serde(tag = "type", rename_all = "lowercase")]
51pub enum LlmInput {
52 Text {
54 prompt: String,
56 },
57 Chat {
59 messages: Vec<ChatMessage>,
61 },
62 Multimodal {
64 parts: Vec<ContentPart>,
66 },
67}
68
69#[derive(Debug, Clone, Serialize, Deserialize)]
71pub struct ChatMessage {
72 pub role: String,
74 pub content: String,
76 #[serde(skip_serializing_if = "Option::is_none")]
78 pub name: Option<String>,
79}
80
81#[derive(Debug, Clone, Serialize, Deserialize)]
83#[serde(tag = "type", rename_all = "lowercase")]
84pub enum ContentPart {
85 Text {
87 text: String,
89 },
90 Image {
92 source: String,
94 },
95 Audio {
97 source: String,
99 },
100}
101
102#[derive(Debug, Clone, Serialize, Deserialize)]
104pub struct LlmOutput {
105 pub content: String,
107 pub finish_reason: Option<String>,
109 #[serde(default)]
111 pub metadata: HashMap<String, serde_json::Value>,
112}
113
114#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
116#[serde(rename_all = "UPPERCASE")]
117pub enum SpanStatus {
118 Ok,
120 Error,
122 Unset,
124}
125
126impl Default for SpanStatus {
127 fn default() -> Self {
128 SpanStatus::Unset
129 }
130}
131
132#[derive(Debug, Clone, Serialize, Deserialize)]
134pub struct SpanEvent {
135 pub name: String,
137 pub timestamp: DateTime<Utc>,
139 #[serde(default)]
141 pub attributes: HashMap<String, serde_json::Value>,
142}
143
144impl LlmSpan {
145 pub fn builder() -> LlmSpanBuilder {
147 LlmSpanBuilder::default()
148 }
149
150 pub fn is_success(&self) -> bool {
152 self.status == SpanStatus::Ok
153 }
154
155 pub fn is_error(&self) -> bool {
157 self.status == SpanStatus::Error
158 }
159
160 pub fn total_tokens(&self) -> Option<u32> {
162 self.token_usage.as_ref().map(|u| u.total_tokens)
163 }
164
165 pub fn total_cost_usd(&self) -> Option<f64> {
167 self.cost.as_ref().map(|c| c.amount_usd)
168 }
169
170 pub fn duration_ms(&self) -> u64 {
172 self.latency.total_ms
173 }
174}
175
176#[derive(Default)]
178pub struct LlmSpanBuilder {
179 span_id: Option<SpanId>,
180 trace_id: Option<TraceId>,
181 parent_span_id: Option<SpanId>,
182 name: Option<String>,
183 provider: Option<Provider>,
184 model: Option<String>,
185 input: Option<LlmInput>,
186 output: Option<LlmOutput>,
187 token_usage: Option<TokenUsage>,
188 cost: Option<Cost>,
189 latency: Option<Latency>,
190 metadata: Option<Metadata>,
191 status: SpanStatus,
192 attributes: HashMap<String, serde_json::Value>,
193 events: Vec<SpanEvent>,
194}
195
196impl LlmSpanBuilder {
197 pub fn span_id(mut self, id: impl Into<SpanId>) -> Self {
199 self.span_id = Some(id.into());
200 self
201 }
202
203 pub fn trace_id(mut self, id: impl Into<TraceId>) -> Self {
205 self.trace_id = Some(id.into());
206 self
207 }
208
209 pub fn parent_span_id(mut self, id: impl Into<SpanId>) -> Self {
211 self.parent_span_id = Some(id.into());
212 self
213 }
214
215 pub fn name(mut self, name: impl Into<String>) -> Self {
217 self.name = Some(name.into());
218 self
219 }
220
221 pub fn provider(mut self, provider: Provider) -> Self {
223 self.provider = Some(provider);
224 self
225 }
226
227 pub fn model(mut self, model: impl Into<String>) -> Self {
229 self.model = Some(model.into());
230 self
231 }
232
233 pub fn input(mut self, input: LlmInput) -> Self {
235 self.input = Some(input);
236 self
237 }
238
239 pub fn output(mut self, output: LlmOutput) -> Self {
241 self.output = Some(output);
242 self
243 }
244
245 pub fn token_usage(mut self, usage: TokenUsage) -> Self {
247 self.token_usage = Some(usage);
248 self
249 }
250
251 pub fn cost(mut self, cost: Cost) -> Self {
253 self.cost = Some(cost);
254 self
255 }
256
257 pub fn latency(mut self, latency: Latency) -> Self {
259 self.latency = Some(latency);
260 self
261 }
262
263 pub fn metadata(mut self, metadata: Metadata) -> Self {
265 self.metadata = Some(metadata);
266 self
267 }
268
269 pub fn status(mut self, status: SpanStatus) -> Self {
271 self.status = status;
272 self
273 }
274
275 pub fn attribute(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
277 self.attributes.insert(key.into(), value);
278 self
279 }
280
281 pub fn event(mut self, event: SpanEvent) -> Self {
283 self.events.push(event);
284 self
285 }
286
287 pub fn build(self) -> Result<LlmSpan, &'static str> {
289 Ok(LlmSpan {
290 span_id: self.span_id.ok_or("span_id is required")?,
291 trace_id: self.trace_id.ok_or("trace_id is required")?,
292 parent_span_id: self.parent_span_id,
293 name: self.name.ok_or("name is required")?,
294 provider: self.provider.ok_or("provider is required")?,
295 model: self.model.ok_or("model is required")?,
296 input: self.input.ok_or("input is required")?,
297 output: self.output,
298 token_usage: self.token_usage,
299 cost: self.cost,
300 latency: self.latency.ok_or("latency is required")?,
301 metadata: self.metadata.unwrap_or_default(),
302 status: self.status,
303 attributes: self.attributes,
304 events: self.events,
305 })
306 }
307}
308
309#[cfg(test)]
310mod tests {
311 use super::*;
312 use chrono::Utc;
313
314 #[test]
315 fn test_span_builder() {
316 let now = Utc::now();
317 let latency = Latency::new(now, now);
318
319 let span = LlmSpan::builder()
320 .span_id("span_123")
321 .trace_id("trace_456")
322 .name("llm.completion")
323 .provider(Provider::OpenAI)
324 .model("gpt-4")
325 .input(LlmInput::Text {
326 prompt: "Hello".to_string(),
327 })
328 .latency(latency)
329 .status(SpanStatus::Ok)
330 .build()
331 .expect("Failed to build span");
332
333 assert_eq!(span.span_id, "span_123");
334 assert_eq!(span.trace_id, "trace_456");
335 assert_eq!(span.provider, Provider::OpenAI);
336 assert!(span.is_success());
337 }
338}