Skip to main content

lc_callbacks/tracing/
span.rs

1use chrono::Utc;
2use serde::{Deserialize, Serialize};
3
4/// Unique span identifier.
5pub type SpanId = String;
6
7/// Kind of span for categorization.
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
9#[serde(rename_all = "snake_case")]
10pub enum SpanKind {
11    /// LLM inference call
12    Llm,
13    /// Chain execution
14    Chain,
15    /// Tool invocation
16    Tool,
17    /// Retriever query
18    Retriever,
19    /// Agent execution
20    Agent,
21    /// Custom span kind
22    Custom(String),
23}
24
25impl std::fmt::Display for SpanKind {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        match self {
28            SpanKind::Llm => write!(f, "llm"),
29            SpanKind::Chain => write!(f, "chain"),
30            SpanKind::Tool => write!(f, "tool"),
31            SpanKind::Retriever => write!(f, "retriever"),
32            SpanKind::Agent => write!(f, "agent"),
33            SpanKind::Custom(name) => write!(f, "custom:{}", name),
34        }
35    }
36}
37
38impl From<crate::RunType> for SpanKind {
39    fn from(run_type: crate::RunType) -> Self {
40        match run_type {
41            crate::RunType::Llm => SpanKind::Llm,
42            crate::RunType::Chain => SpanKind::Chain,
43            crate::RunType::Tool => SpanKind::Tool,
44            crate::RunType::Retriever => SpanKind::Retriever,
45            // RunType variants without a dedicated SpanKind collapse to Custom
46            crate::RunType::Embedding => SpanKind::Custom("embedding".to_string()),
47            crate::RunType::Prompt => SpanKind::Custom("prompt".to_string()),
48            crate::RunType::Parser => SpanKind::Custom("parser".to_string()),
49        }
50    }
51}
52
53/// Token usage recorded in a span.
54#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
55pub struct SpanTokenUsage {
56    /// Number of tokens in the prompt.
57    pub prompt_tokens: usize,
58    /// Number of tokens in the completion.
59    pub completion_tokens: usize,
60    /// Total number of tokens.
61    pub total_tokens: usize,
62    /// `gen_ai.usage.cache_read.input_tokens` (Development extension).
63    #[serde(default, skip_serializing_if = "Option::is_none")]
64    pub cache_read_input_tokens: Option<usize>,
65    /// `gen_ai.usage.cache_write.input_tokens` (Development extension; 2026
66    /// rename of `cache_creation`).
67    #[serde(default, skip_serializing_if = "Option::is_none")]
68    pub cache_write_input_tokens: Option<usize>,
69    /// `gen_ai.usage.reasoning.output_tokens` (Development extension).
70    #[serde(default, skip_serializing_if = "Option::is_none")]
71    pub reasoning_output_tokens: Option<usize>,
72}
73
74/// Status of a span.
75#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
76#[serde(rename_all = "snake_case")]
77pub enum SpanStatus {
78    /// Span completed successfully
79    Ok,
80    /// Span ended with an error
81    Error(String),
82}
83
84/// A single trace span with parent-child relationships.
85#[derive(Debug, Clone, Serialize, Deserialize)]
86pub struct TraceSpan {
87    /// Unique span identifier
88    pub id: SpanId,
89    /// Parent span ID (None for root spans)
90    pub parent_id: Option<SpanId>,
91    /// Human-readable span name
92    pub name: String,
93    /// Span category
94    pub kind: SpanKind,
95    /// ISO 8601 start time
96    pub start_time: Option<String>,
97    /// ISO 8601 end time
98    pub end_time: Option<String>,
99    /// Token usage (for LLM spans)
100    pub tokens: Option<SpanTokenUsage>,
101    /// Estimated cost in USD
102    pub cost: Option<f64>,
103    /// Measured latency in milliseconds
104    pub latency_ms: Option<u64>,
105    /// Arbitrary key-value metadata
106    pub metadata: serde_json::Value,
107    /// Span completion status
108    pub status: SpanStatus,
109
110    // --- OTel GenAI SemConv fields ---
111    /// gen_ai.system: The LLM provider name (e.g., "openai", "anthropic")
112    #[serde(skip_serializing_if = "Option::is_none")]
113    pub gen_ai_system: Option<String>,
114    /// gen_ai.request.model: The model requested
115    #[serde(skip_serializing_if = "Option::is_none")]
116    pub gen_ai_request_model: Option<String>,
117    /// gen_ai.response.model: The actual model used
118    #[serde(skip_serializing_if = "Option::is_none")]
119    pub gen_ai_response_model: Option<String>,
120    /// gen_ai.response.finish_reason: Why the model stopped generating
121    #[serde(skip_serializing_if = "Option::is_none")]
122    pub gen_ai_finish_reason: Option<String>,
123    /// gen_ai.request.max_tokens: Maximum tokens requested
124    #[serde(skip_serializing_if = "Option::is_none")]
125    pub gen_ai_request_max_tokens: Option<u64>,
126    /// gen_ai.request.temperature: Temperature parameter
127    #[serde(skip_serializing_if = "Option::is_none")]
128    pub gen_ai_request_temperature: Option<f64>,
129    /// gen_ai.operation.name: The operation name (chat, completion)
130    #[serde(skip_serializing_if = "Option::is_none")]
131    pub gen_ai_operation_name: Option<String>,
132    /// gen_ai.tool.name: The tool name (for tool spans)
133    #[serde(skip_serializing_if = "Option::is_none")]
134    pub gen_ai_tool_name: Option<String>,
135}
136
137/// A node in the trace tree (span + children).
138#[derive(Debug, Clone, Serialize, Deserialize)]
139pub struct TraceNode {
140    /// The span itself.
141    pub span: TraceSpan,
142    /// Child nodes of this span.
143    pub children: Vec<TraceNode>,
144}
145
146/// Estimates the USD cost of a single span from its token usage and gen_ai model.
147///
148/// Returns `None` when the span has no tokens or no recognizable model.
149fn estimate_span_cost(span: &TraceSpan) -> Option<f64> {
150    let tokens = span.tokens.as_ref()?;
151    let model = span
152        .gen_ai_request_model
153        .as_deref()
154        .or(span.gen_ai_response_model.as_deref())?;
155    crate::pricing::estimate_cost_usd(tokens.prompt_tokens, tokens.completion_tokens, model)
156}
157
158/// Sums the USD cost across spans (E3 agent roll-up).
159///
160/// Uses each span's recorded `cost` when set; otherwise estimates it from token
161/// usage + gen_ai model. Unknown models contribute zero (an unestimated span can't
162/// be trusted, so it doesn't deflate the aggregate with a guessed value).
163pub fn aggregate_cost(spans: &[TraceSpan]) -> f64 {
164    spans
165        .iter()
166        .map(|s| s.cost.or_else(|| estimate_span_cost(s)).unwrap_or(0.0))
167        .sum()
168}
169
170pub(crate) fn build_tree(root: &TraceSpan, all_spans: &[TraceSpan]) -> TraceNode {
171    let children: Vec<TraceNode> = all_spans
172        .iter()
173        .filter(|s| s.parent_id.as_deref() == Some(root.id.as_str()))
174        .map(|child| build_tree(child, all_spans))
175        .collect();
176
177    TraceNode {
178        span: root.clone(),
179        children,
180    }
181}
182
183/// Helper to create a new span with common defaults.
184pub(crate) fn make_span(
185    id: String,
186    parent_id: Option<SpanId>,
187    name: &str,
188    kind: SpanKind,
189) -> TraceSpan {
190    TraceSpan {
191        id,
192        parent_id,
193        name: name.to_string(),
194        kind,
195        start_time: Some(Utc::now().to_rfc3339()),
196        end_time: None,
197        tokens: None,
198        cost: None,
199        latency_ms: None,
200        metadata: serde_json::Value::Object(serde_json::Map::new()),
201        status: SpanStatus::Ok,
202        gen_ai_system: None,
203        gen_ai_request_model: None,
204        gen_ai_response_model: None,
205        gen_ai_finish_reason: None,
206        gen_ai_request_max_tokens: None,
207        gen_ai_request_temperature: None,
208        gen_ai_operation_name: None,
209        gen_ai_tool_name: None,
210    }
211}