lc_callbacks/tracing/
span.rs1use chrono::Utc;
2use serde::{Deserialize, Serialize};
3
4pub type SpanId = String;
6
7#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
9#[serde(rename_all = "snake_case")]
10pub enum SpanKind {
11 Llm,
13 Chain,
15 Tool,
17 Retriever,
19 Agent,
21 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#[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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
48#[serde(rename_all = "snake_case")]
49pub enum SpanStatus {
50 Ok,
52 Error(String),
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct TraceSpan {
59 pub id: SpanId,
61 pub parent_id: Option<SpanId>,
63 pub name: String,
65 pub kind: SpanKind,
67 pub start_time: Option<String>,
69 pub end_time: Option<String>,
71 pub tokens: Option<SpanTokenUsage>,
73 pub cost: Option<f64>,
75 pub latency_ms: Option<u64>,
77 pub metadata: serde_json::Value,
79 pub status: SpanStatus,
81}
82
83#[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
103pub(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}