llm_observatory_core/
span.rs

1// Copyright 2025 LLM Observatory Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4//! LLM span definitions following OpenTelemetry GenAI semantic conventions.
5
6use crate::types::{Cost, Latency, Metadata, Provider, TokenUsage, TraceId, SpanId};
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10
11/// Represents a single LLM operation (request/response) as an OpenTelemetry span.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct LlmSpan {
14    /// Unique span identifier
15    pub span_id: SpanId,
16    /// Trace identifier this span belongs to
17    pub trace_id: TraceId,
18    /// Parent span identifier (if part of a chain)
19    pub parent_span_id: Option<SpanId>,
20    /// Span name/operation type
21    pub name: String,
22    /// LLM provider
23    pub provider: Provider,
24    /// Model name
25    pub model: String,
26    /// Request input (prompt)
27    pub input: LlmInput,
28    /// Response output
29    pub output: Option<LlmOutput>,
30    /// Token usage statistics
31    pub token_usage: Option<TokenUsage>,
32    /// Cost information
33    pub cost: Option<Cost>,
34    /// Latency metrics
35    pub latency: Latency,
36    /// Metadata and tags
37    pub metadata: Metadata,
38    /// Span status
39    pub status: SpanStatus,
40    /// OpenTelemetry attributes
41    #[serde(default)]
42    pub attributes: HashMap<String, serde_json::Value>,
43    /// Events recorded during span
44    #[serde(default)]
45    pub events: Vec<SpanEvent>,
46}
47
48/// LLM input (prompt).
49#[derive(Debug, Clone, Serialize, Deserialize)]
50#[serde(tag = "type", rename_all = "lowercase")]
51pub enum LlmInput {
52    /// Simple text prompt
53    Text {
54        /// The prompt text
55        prompt: String,
56    },
57    /// Chat messages
58    Chat {
59        /// Array of messages
60        messages: Vec<ChatMessage>,
61    },
62    /// Multimodal input
63    Multimodal {
64        /// Content parts
65        parts: Vec<ContentPart>,
66    },
67}
68
69/// Chat message for conversational models.
70#[derive(Debug, Clone, Serialize, Deserialize)]
71pub struct ChatMessage {
72    /// Role (system, user, assistant)
73    pub role: String,
74    /// Message content
75    pub content: String,
76    /// Optional message name
77    #[serde(skip_serializing_if = "Option::is_none")]
78    pub name: Option<String>,
79}
80
81/// Content part for multimodal inputs.
82#[derive(Debug, Clone, Serialize, Deserialize)]
83#[serde(tag = "type", rename_all = "lowercase")]
84pub enum ContentPart {
85    /// Text content
86    Text {
87        /// The text
88        text: String,
89    },
90    /// Image content
91    Image {
92        /// Image URL or base64 data
93        source: String,
94    },
95    /// Audio content
96    Audio {
97        /// Audio URL or base64 data
98        source: String,
99    },
100}
101
102/// LLM output (completion).
103#[derive(Debug, Clone, Serialize, Deserialize)]
104pub struct LlmOutput {
105    /// Generated text
106    pub content: String,
107    /// Finish reason (stop, length, content_filter, etc.)
108    pub finish_reason: Option<String>,
109    /// Additional output metadata
110    #[serde(default)]
111    pub metadata: HashMap<String, serde_json::Value>,
112}
113
114/// Span status following OpenTelemetry conventions.
115#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
116#[serde(rename_all = "UPPERCASE")]
117pub enum SpanStatus {
118    /// Operation completed successfully
119    Ok,
120    /// Operation failed
121    Error,
122    /// Status not set
123    Unset,
124}
125
126impl Default for SpanStatus {
127    fn default() -> Self {
128        SpanStatus::Unset
129    }
130}
131
132/// Event recorded during span execution.
133#[derive(Debug, Clone, Serialize, Deserialize)]
134pub struct SpanEvent {
135    /// Event name
136    pub name: String,
137    /// Timestamp when event occurred
138    pub timestamp: DateTime<Utc>,
139    /// Event attributes
140    #[serde(default)]
141    pub attributes: HashMap<String, serde_json::Value>,
142}
143
144impl LlmSpan {
145    /// Create a new LLM span builder.
146    pub fn builder() -> LlmSpanBuilder {
147        LlmSpanBuilder::default()
148    }
149
150    /// Check if the span represents a successful operation.
151    pub fn is_success(&self) -> bool {
152        self.status == SpanStatus::Ok
153    }
154
155    /// Check if the span represents a failed operation.
156    pub fn is_error(&self) -> bool {
157        self.status == SpanStatus::Error
158    }
159
160    /// Get total tokens used (if available).
161    pub fn total_tokens(&self) -> Option<u32> {
162        self.token_usage.as_ref().map(|u| u.total_tokens)
163    }
164
165    /// Get total cost in USD (if available).
166    pub fn total_cost_usd(&self) -> Option<f64> {
167        self.cost.as_ref().map(|c| c.amount_usd)
168    }
169
170    /// Get duration in milliseconds.
171    pub fn duration_ms(&self) -> u64 {
172        self.latency.total_ms
173    }
174}
175
176/// Builder for creating LlmSpan instances.
177#[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    /// Set span ID.
198    pub fn span_id(mut self, id: impl Into<SpanId>) -> Self {
199        self.span_id = Some(id.into());
200        self
201    }
202
203    /// Set trace ID.
204    pub fn trace_id(mut self, id: impl Into<TraceId>) -> Self {
205        self.trace_id = Some(id.into());
206        self
207    }
208
209    /// Set parent span ID.
210    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    /// Set span name.
216    pub fn name(mut self, name: impl Into<String>) -> Self {
217        self.name = Some(name.into());
218        self
219    }
220
221    /// Set provider.
222    pub fn provider(mut self, provider: Provider) -> Self {
223        self.provider = Some(provider);
224        self
225    }
226
227    /// Set model name.
228    pub fn model(mut self, model: impl Into<String>) -> Self {
229        self.model = Some(model.into());
230        self
231    }
232
233    /// Set input.
234    pub fn input(mut self, input: LlmInput) -> Self {
235        self.input = Some(input);
236        self
237    }
238
239    /// Set output.
240    pub fn output(mut self, output: LlmOutput) -> Self {
241        self.output = Some(output);
242        self
243    }
244
245    /// Set token usage.
246    pub fn token_usage(mut self, usage: TokenUsage) -> Self {
247        self.token_usage = Some(usage);
248        self
249    }
250
251    /// Set cost.
252    pub fn cost(mut self, cost: Cost) -> Self {
253        self.cost = Some(cost);
254        self
255    }
256
257    /// Set latency.
258    pub fn latency(mut self, latency: Latency) -> Self {
259        self.latency = Some(latency);
260        self
261    }
262
263    /// Set metadata.
264    pub fn metadata(mut self, metadata: Metadata) -> Self {
265        self.metadata = Some(metadata);
266        self
267    }
268
269    /// Set status.
270    pub fn status(mut self, status: SpanStatus) -> Self {
271        self.status = status;
272        self
273    }
274
275    /// Add an attribute.
276    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    /// Add an event.
282    pub fn event(mut self, event: SpanEvent) -> Self {
283        self.events.push(event);
284        self
285    }
286
287    /// Build the LlmSpan.
288    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}