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
38/// Token usage recorded in a span.
39#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
40pub struct SpanTokenUsage {
41    pub prompt_tokens: usize,
42    pub completion_tokens: usize,
43    pub total_tokens: usize,
44}
45
46/// Status of a span.
47#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
48#[serde(rename_all = "snake_case")]
49pub enum SpanStatus {
50    /// Span completed successfully
51    Ok,
52    /// Span ended with an error
53    Error(String),
54}
55
56/// A single trace span with parent-child relationships.
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct TraceSpan {
59    /// Unique span identifier
60    pub id: SpanId,
61    /// Parent span ID (None for root spans)
62    pub parent_id: Option<SpanId>,
63    /// Human-readable span name
64    pub name: String,
65    /// Span category
66    pub kind: SpanKind,
67    /// ISO 8601 start time
68    pub start_time: Option<String>,
69    /// ISO 8601 end time
70    pub end_time: Option<String>,
71    /// Token usage (for LLM spans)
72    pub tokens: Option<SpanTokenUsage>,
73    /// Estimated cost in USD
74    pub cost: Option<f64>,
75    /// Measured latency in milliseconds
76    pub latency_ms: Option<u64>,
77    /// Arbitrary key-value metadata
78    pub metadata: serde_json::Value,
79    /// Span completion status
80    pub status: SpanStatus,
81}
82
83/// A node in the trace tree (span + children).
84#[derive(Debug, Clone, Serialize, Deserialize)]
85pub struct TraceNode {
86    pub span: TraceSpan,
87    pub children: Vec<TraceNode>,
88}
89
90pub(crate) fn build_tree(root: &TraceSpan, all_spans: &[TraceSpan]) -> TraceNode {
91    let children: Vec<TraceNode> = all_spans
92        .iter()
93        .filter(|s| s.parent_id.as_deref() == Some(root.id.as_str()))
94        .map(|child| build_tree(child, all_spans))
95        .collect();
96
97    TraceNode {
98        span: root.clone(),
99        children,
100    }
101}
102
103/// Helper to create a new span with common defaults.
104pub(crate) fn make_span(
105    id: String,
106    parent_id: Option<SpanId>,
107    name: &str,
108    kind: SpanKind,
109) -> TraceSpan {
110    TraceSpan {
111        id,
112        parent_id,
113        name: name.to_string(),
114        kind,
115        start_time: Some(Utc::now().to_rfc3339()),
116        end_time: None,
117        tokens: None,
118        cost: None,
119        latency_ms: None,
120        metadata: serde_json::Value::Object(serde_json::Map::new()),
121        status: SpanStatus::Ok,
122    }
123}