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
136fn estimate_span_cost(span: &TraceSpan) -> Option<f64> {
140 let tokens = span.tokens.as_ref()?;
141 let model = span
142 .gen_ai_request_model
143 .as_deref()
144 .or(span.gen_ai_response_model.as_deref())?;
145 crate::pricing::estimate_cost_usd(tokens.prompt_tokens, tokens.completion_tokens, model)
146}
147
148pub fn aggregate_cost(spans: &[TraceSpan]) -> f64 {
154 spans
155 .iter()
156 .map(|s| s.cost.or_else(|| estimate_span_cost(s)).unwrap_or(0.0))
157 .sum()
158}
159
160pub(crate) fn build_tree(root: &TraceSpan, all_spans: &[TraceSpan]) -> TraceNode {
161 let children: Vec<TraceNode> = all_spans
162 .iter()
163 .filter(|s| s.parent_id.as_deref() == Some(root.id.as_str()))
164 .map(|child| build_tree(child, all_spans))
165 .collect();
166
167 TraceNode {
168 span: root.clone(),
169 children,
170 }
171}
172
173pub(crate) fn make_span(
175 id: String,
176 parent_id: Option<SpanId>,
177 name: &str,
178 kind: SpanKind,
179) -> TraceSpan {
180 TraceSpan {
181 id,
182 parent_id,
183 name: name.to_string(),
184 kind,
185 start_time: Some(Utc::now().to_rfc3339()),
186 end_time: None,
187 tokens: None,
188 cost: None,
189 latency_ms: None,
190 metadata: serde_json::Value::Object(serde_json::Map::new()),
191 status: SpanStatus::Ok,
192 gen_ai_system: None,
193 gen_ai_request_model: None,
194 gen_ai_response_model: None,
195 gen_ai_finish_reason: None,
196 gen_ai_request_max_tokens: None,
197 gen_ai_request_temperature: None,
198 gen_ai_operation_name: None,
199 gen_ai_tool_name: None,
200 }
201}