1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use serde_json::Value as JsonValue;
4use std::collections::HashMap;
5
6#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
8#[serde(rename_all = "snake_case")]
9pub enum TraceStatus {
10 Running,
11 Success,
12 Error,
13 Interrupted,
14}
15
16#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
18#[serde(rename_all = "snake_case")]
19pub enum SpanType {
20 GraphNode,
22 LlmGeneration,
24 ToolCall,
26}
27
28#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
30#[serde(rename_all = "snake_case")]
31pub enum SpanStatus {
32 Running,
33 Success,
34 Error,
35}
36
37#[derive(Debug, Clone, Serialize, Deserialize)]
39pub struct Trace {
40 pub id: String,
41 pub name: String,
42 pub input: JsonValue,
43 pub output: Option<JsonValue>,
44 pub status: TraceStatus,
45 pub start_time: DateTime<Utc>,
46 pub end_time: Option<DateTime<Utc>>,
47 pub metadata: HashMap<String, JsonValue>,
48}
49
50impl Trace {
51 pub fn new(id: String, name: String, input: JsonValue) -> Self {
52 Self {
53 id,
54 name,
55 input,
56 output: None,
57 status: TraceStatus::Running,
58 start_time: Utc::now(),
59 end_time: None,
60 metadata: HashMap::new(),
61 }
62 }
63
64 pub fn duration_ms(&self) -> Option<u64> {
65 let end = self.end_time.unwrap_or_else(Utc::now);
66 let dur = (end - self.start_time).num_milliseconds();
67 if dur < 0 { Some(0) } else { Some(dur as u64) }
68 }
69
70 pub fn finish(&mut self, output: JsonValue, status: TraceStatus) {
71 self.output = Some(output);
72 self.status = status;
73 self.end_time = Some(Utc::now());
74 }
75}
76
77#[derive(Debug, Clone, Serialize, Deserialize)]
79pub struct Span {
80 pub id: String,
81 pub trace_id: String,
82 pub parent_span_id: Option<String>,
83 pub name: String,
84 pub span_type: SpanType,
85 pub input: JsonValue,
86 pub output: Option<JsonValue>,
87 pub status: SpanStatus,
88 pub start_time: DateTime<Utc>,
89 pub end_time: Option<DateTime<Utc>>,
90 pub metadata: SpanMetadata,
92}
93
94impl Span {
95 pub fn new(
96 id: String,
97 trace_id: String,
98 parent_span_id: Option<String>,
99 name: String,
100 span_type: SpanType,
101 input: JsonValue,
102 ) -> Self {
103 Self {
104 id,
105 trace_id,
106 parent_span_id,
107 name,
108 span_type,
109 input,
110 output: None,
111 status: SpanStatus::Running,
112 start_time: Utc::now(),
113 end_time: None,
114 metadata: SpanMetadata::default(),
115 }
116 }
117
118 pub fn duration_ms(&self) -> Option<u64> {
119 let end = self.end_time.unwrap_or_else(Utc::now);
120 let dur = (end - self.start_time).num_milliseconds();
121 if dur < 0 { Some(0) } else { Some(dur as u64) }
122 }
123
124 pub fn finish(&mut self, output: JsonValue, status: SpanStatus) {
125 self.output = Some(output);
126 self.status = status;
127 self.end_time = Some(Utc::now());
128 }
129}
130
131#[derive(Debug, Clone, Default, Serialize, Deserialize)]
133pub struct SpanMetadata {
134 pub model: Option<String>,
136 pub provider: Option<String>,
138 pub tokens_in: Option<u32>,
140 pub tokens_out: Option<u32>,
142 pub total_tokens: Option<u32>,
144 pub tool_name: Option<String>,
146 pub extra: HashMap<String, JsonValue>,
148}
149
150#[derive(Debug, Clone, Serialize, Deserialize)]
152pub struct TraceSummary {
153 pub id: String,
154 pub name: String,
155 pub status: TraceStatus,
156 pub start_time: DateTime<Utc>,
157 pub end_time: Option<DateTime<Utc>>,
158 pub duration_ms: Option<u64>,
159 pub span_count: usize,
160}
161
162impl From<&Trace> for TraceSummary {
163 fn from(trace: &Trace) -> Self {
164 Self {
165 id: trace.id.clone(),
166 name: trace.name.clone(),
167 status: trace.status,
168 start_time: trace.start_time,
169 end_time: trace.end_time,
170 duration_ms: trace.duration_ms(),
171 span_count: 0,
172 }
173 }
174}
175
176#[derive(Debug, Clone, Serialize, Deserialize)]
178pub struct TraceDetail {
179 pub trace: Trace,
180 pub spans: Vec<Span>,
181}