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
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 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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
55pub struct SpanTokenUsage {
56 pub prompt_tokens: usize,
58 pub completion_tokens: usize,
60 pub total_tokens: usize,
62}
63
64#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
66#[serde(rename_all = "snake_case")]
67pub enum SpanStatus {
68 Ok,
70 Error(String),
72}
73
74#[derive(Debug, Clone, Serialize, Deserialize)]
76pub struct TraceSpan {
77 pub id: SpanId,
79 pub parent_id: Option<SpanId>,
81 pub name: String,
83 pub kind: SpanKind,
85 pub start_time: Option<String>,
87 pub end_time: Option<String>,
89 pub tokens: Option<SpanTokenUsage>,
91 pub cost: Option<f64>,
93 pub latency_ms: Option<u64>,
95 pub metadata: serde_json::Value,
97 pub status: SpanStatus,
99
100 #[serde(skip_serializing_if = "Option::is_none")]
103 pub gen_ai_system: Option<String>,
104 #[serde(skip_serializing_if = "Option::is_none")]
106 pub gen_ai_request_model: Option<String>,
107 #[serde(skip_serializing_if = "Option::is_none")]
109 pub gen_ai_response_model: Option<String>,
110 #[serde(skip_serializing_if = "Option::is_none")]
112 pub gen_ai_finish_reason: Option<String>,
113 #[serde(skip_serializing_if = "Option::is_none")]
115 pub gen_ai_request_max_tokens: Option<u64>,
116 #[serde(skip_serializing_if = "Option::is_none")]
118 pub gen_ai_request_temperature: Option<f64>,
119 #[serde(skip_serializing_if = "Option::is_none")]
121 pub gen_ai_operation_name: Option<String>,
122 #[serde(skip_serializing_if = "Option::is_none")]
124 pub gen_ai_tool_name: Option<String>,
125}
126
127#[derive(Debug, Clone, Serialize, Deserialize)]
129pub struct TraceNode {
130 pub span: TraceSpan,
132 pub children: Vec<TraceNode>,
134}
135
136pub(crate) fn build_tree(root: &TraceSpan, all_spans: &[TraceSpan]) -> TraceNode {
137 let children: Vec<TraceNode> = all_spans
138 .iter()
139 .filter(|s| s.parent_id.as_deref() == Some(root.id.as_str()))
140 .map(|child| build_tree(child, all_spans))
141 .collect();
142
143 TraceNode {
144 span: root.clone(),
145 children,
146 }
147}
148
149pub(crate) fn make_span(
151 id: String,
152 parent_id: Option<SpanId>,
153 name: &str,
154 kind: SpanKind,
155) -> TraceSpan {
156 TraceSpan {
157 id,
158 parent_id,
159 name: name.to_string(),
160 kind,
161 start_time: Some(Utc::now().to_rfc3339()),
162 end_time: None,
163 tokens: None,
164 cost: None,
165 latency_ms: None,
166 metadata: serde_json::Value::Object(serde_json::Map::new()),
167 status: SpanStatus::Ok,
168 gen_ai_system: None,
169 gen_ai_request_model: None,
170 gen_ai_response_model: None,
171 gen_ai_finish_reason: None,
172 gen_ai_request_max_tokens: None,
173 gen_ai_request_temperature: None,
174 gen_ai_operation_name: None,
175 gen_ai_tool_name: None,
176 }
177}