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, Default, Serialize, Deserialize, PartialEq)]
55pub struct SpanTokenUsage {
56 pub prompt_tokens: usize,
58 pub completion_tokens: usize,
60 pub total_tokens: usize,
62 #[serde(default, skip_serializing_if = "Option::is_none")]
64 pub cache_read_input_tokens: Option<usize>,
65 #[serde(default, skip_serializing_if = "Option::is_none")]
68 pub cache_write_input_tokens: Option<usize>,
69 #[serde(default, skip_serializing_if = "Option::is_none")]
71 pub reasoning_output_tokens: Option<usize>,
72}
73
74#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
76#[serde(rename_all = "snake_case")]
77pub enum SpanStatus {
78 Ok,
80 Error(String),
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize)]
86pub struct TraceSpan {
87 pub id: SpanId,
89 pub parent_id: Option<SpanId>,
91 pub name: String,
93 pub kind: SpanKind,
95 pub start_time: Option<String>,
97 pub end_time: Option<String>,
99 pub tokens: Option<SpanTokenUsage>,
101 pub cost: Option<f64>,
103 pub latency_ms: Option<u64>,
105 pub metadata: serde_json::Value,
107 pub status: SpanStatus,
109
110 #[serde(skip_serializing_if = "Option::is_none")]
113 pub gen_ai_system: Option<String>,
114 #[serde(skip_serializing_if = "Option::is_none")]
116 pub gen_ai_request_model: Option<String>,
117 #[serde(skip_serializing_if = "Option::is_none")]
119 pub gen_ai_response_model: Option<String>,
120 #[serde(skip_serializing_if = "Option::is_none")]
122 pub gen_ai_finish_reason: Option<String>,
123 #[serde(skip_serializing_if = "Option::is_none")]
125 pub gen_ai_request_max_tokens: Option<u64>,
126 #[serde(skip_serializing_if = "Option::is_none")]
128 pub gen_ai_request_temperature: Option<f64>,
129 #[serde(skip_serializing_if = "Option::is_none")]
131 pub gen_ai_operation_name: Option<String>,
132 #[serde(skip_serializing_if = "Option::is_none")]
134 pub gen_ai_tool_name: Option<String>,
135}
136
137#[derive(Debug, Clone, Serialize, Deserialize)]
139pub struct TraceNode {
140 pub span: TraceSpan,
142 pub children: Vec<TraceNode>,
144}
145
146fn estimate_span_cost(span: &TraceSpan) -> Option<f64> {
150 let tokens = span.tokens.as_ref()?;
151 let model = span
152 .gen_ai_request_model
153 .as_deref()
154 .or(span.gen_ai_response_model.as_deref())?;
155 crate::pricing::estimate_cost_usd(tokens.prompt_tokens, tokens.completion_tokens, model)
156}
157
158pub fn aggregate_cost(spans: &[TraceSpan]) -> f64 {
164 spans
165 .iter()
166 .map(|s| s.cost.or_else(|| estimate_span_cost(s)).unwrap_or(0.0))
167 .sum()
168}
169
170pub(crate) fn build_tree(root: &TraceSpan, all_spans: &[TraceSpan]) -> TraceNode {
171 let children: Vec<TraceNode> = all_spans
172 .iter()
173 .filter(|s| s.parent_id.as_deref() == Some(root.id.as_str()))
174 .map(|child| build_tree(child, all_spans))
175 .collect();
176
177 TraceNode {
178 span: root.clone(),
179 children,
180 }
181}
182
183pub(crate) fn make_span(
185 id: String,
186 parent_id: Option<SpanId>,
187 name: &str,
188 kind: SpanKind,
189) -> TraceSpan {
190 TraceSpan {
191 id,
192 parent_id,
193 name: name.to_string(),
194 kind,
195 start_time: Some(Utc::now().to_rfc3339()),
196 end_time: None,
197 tokens: None,
198 cost: None,
199 latency_ms: None,
200 metadata: serde_json::Value::Object(serde_json::Map::new()),
201 status: SpanStatus::Ok,
202 gen_ai_system: None,
203 gen_ai_request_model: None,
204 gen_ai_response_model: None,
205 gen_ai_finish_reason: None,
206 gen_ai_request_max_tokens: None,
207 gen_ai_request_temperature: None,
208 gen_ai_operation_name: None,
209 gen_ai_tool_name: None,
210 }
211}