1use 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 #[serde(skip_serializing_if = "Option::is_none")]
85 pub gen_ai_system: Option<String>,
86 #[serde(skip_serializing_if = "Option::is_none")]
88 pub gen_ai_request_model: Option<String>,
89 #[serde(skip_serializing_if = "Option::is_none")]
91 pub gen_ai_response_model: Option<String>,
92 #[serde(skip_serializing_if = "Option::is_none")]
94 pub gen_ai_finish_reason: Option<String>,
95 #[serde(skip_serializing_if = "Option::is_none")]
97 pub gen_ai_request_max_tokens: Option<u64>,
98 #[serde(skip_serializing_if = "Option::is_none")]
100 pub gen_ai_request_temperature: Option<f64>,
101 #[serde(skip_serializing_if = "Option::is_none")]
103 pub gen_ai_operation_name: Option<String>,
104 #[serde(skip_serializing_if = "Option::is_none")]
106 pub gen_ai_tool_name: Option<String>,
107}
108
109#[derive(Debug, Clone, Serialize, Deserialize)]
111pub struct TraceNode {
112 pub span: TraceSpan,
113 pub children: Vec<TraceNode>,
114}
115
116pub(crate) fn build_tree(root: &TraceSpan, all_spans: &[TraceSpan]) -> TraceNode {
117 let children: Vec<TraceNode> = all_spans
118 .iter()
119 .filter(|s| s.parent_id.as_deref() == Some(root.id.as_str()))
120 .map(|child| build_tree(child, all_spans))
121 .collect();
122
123 TraceNode {
124 span: root.clone(),
125 children,
126 }
127}
128
129pub(crate) fn make_span(
131 id: String,
132 parent_id: Option<SpanId>,
133 name: &str,
134 kind: SpanKind,
135) -> TraceSpan {
136 TraceSpan {
137 id,
138 parent_id,
139 name: name.to_string(),
140 kind,
141 start_time: Some(Utc::now().to_rfc3339()),
142 end_time: None,
143 tokens: None,
144 cost: None,
145 latency_ms: None,
146 metadata: serde_json::Value::Object(serde_json::Map::new()),
147 status: SpanStatus::Ok,
148 gen_ai_system: None,
149 gen_ai_request_model: None,
150 gen_ai_response_model: None,
151 gen_ai_finish_reason: None,
152 gen_ai_request_max_tokens: None,
153 gen_ai_request_temperature: None,
154 gen_ai_operation_name: None,
155 gen_ai_tool_name: None,
156 }
157}