Skip to main content

synth_ai_core/tracing/
models.rs

1//! Tracing data models.
2//!
3//! This module contains all the data structures for trace recording,
4//! corresponding to Python's `synth_ai.data.traces` and `synth_ai.data.llm_calls`.
5
6use chrono::{DateTime, Utc};
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::collections::HashMap;
10
11// ============================================================================
12// TIME & CONTENT
13// ============================================================================
14
15/// Time record for events and messages.
16#[derive(Debug, Clone, Default, Serialize, Deserialize)]
17pub struct TimeRecord {
18    /// Unix timestamp for the event
19    #[serde(default)]
20    pub event_time: f64,
21    /// Optional message-specific timestamp
22    #[serde(default)]
23    pub message_time: Option<i64>,
24}
25
26impl TimeRecord {
27    /// Create a new time record with the current time.
28    pub fn now() -> Self {
29        Self {
30            event_time: Utc::now().timestamp_millis() as f64 / 1000.0,
31            message_time: None,
32        }
33    }
34}
35
36/// Content for messages, supporting text or JSON.
37#[derive(Debug, Clone, Default, Serialize, Deserialize)]
38pub struct MessageContent {
39    /// Plain text content
40    #[serde(default)]
41    pub text: Option<String>,
42    /// JSON-serialized content
43    #[serde(default)]
44    pub json_payload: Option<String>,
45}
46
47impl MessageContent {
48    /// Create from text.
49    pub fn from_text(text: impl Into<String>) -> Self {
50        Self {
51            text: Some(text.into()),
52            json_payload: None,
53        }
54    }
55
56    /// Create from JSON value.
57    pub fn from_json(value: &Value) -> Self {
58        Self {
59            text: None,
60            json_payload: Some(value.to_string()),
61        }
62    }
63
64    /// Get content as text (either raw text or JSON string).
65    pub fn as_text(&self) -> Option<&str> {
66        self.text.as_deref().or(self.json_payload.as_deref())
67    }
68}
69
70// ============================================================================
71// EVENTS
72// ============================================================================
73
74/// Event type discriminator.
75#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
76#[serde(rename_all = "snake_case")]
77pub enum EventType {
78    /// LLM/CAIS call event
79    Cais,
80    /// Environment step event (Gym-style)
81    Environment,
82    /// Runtime/action selection event
83    Runtime,
84}
85
86impl std::fmt::Display for EventType {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        match self {
89            EventType::Cais => write!(f, "cais"),
90            EventType::Environment => write!(f, "environment"),
91            EventType::Runtime => write!(f, "runtime"),
92        }
93    }
94}
95
96/// Base fields common to all events.
97#[derive(Debug, Clone, Default, Serialize, Deserialize)]
98pub struct BaseEventFields {
99    /// System/component instance ID
100    #[serde(default)]
101    pub system_instance_id: String,
102    /// Time record
103    #[serde(default)]
104    pub time_record: TimeRecord,
105    /// Event metadata (key-value pairs)
106    #[serde(default)]
107    pub metadata: HashMap<String, Value>,
108    /// Structured event metadata
109    #[serde(default)]
110    pub event_metadata: Option<Vec<Value>>,
111}
112
113impl BaseEventFields {
114    /// Create new base fields with a system instance ID.
115    pub fn new(system_instance_id: impl Into<String>) -> Self {
116        Self {
117            system_instance_id: system_instance_id.into(),
118            time_record: TimeRecord::now(),
119            metadata: HashMap::new(),
120            event_metadata: None,
121        }
122    }
123}
124
125/// LLM/CAIS call event.
126#[derive(Debug, Clone, Default, Serialize, Deserialize)]
127pub struct LMCAISEvent {
128    /// Base event fields
129    #[serde(flatten)]
130    pub base: BaseEventFields,
131    /// Model name (e.g., "gpt-4", "claude-3-opus")
132    #[serde(default)]
133    pub model_name: String,
134    /// Provider (e.g., "openai", "anthropic")
135    #[serde(default)]
136    pub provider: Option<String>,
137    /// Input/prompt tokens
138    #[serde(default)]
139    pub input_tokens: Option<i32>,
140    /// Output/completion tokens
141    #[serde(default)]
142    pub output_tokens: Option<i32>,
143    /// Total tokens
144    #[serde(default)]
145    pub total_tokens: Option<i32>,
146    /// Cost in USD
147    #[serde(default)]
148    pub cost_usd: Option<f64>,
149    /// Latency in milliseconds
150    #[serde(default)]
151    pub latency_ms: Option<i32>,
152    /// OpenTelemetry span ID
153    #[serde(default)]
154    pub span_id: Option<String>,
155    /// OpenTelemetry trace ID
156    #[serde(default)]
157    pub trace_id: Option<String>,
158    /// Detailed call records
159    #[serde(default)]
160    pub call_records: Option<Vec<LLMCallRecord>>,
161    /// System state before the call
162    #[serde(default)]
163    pub system_state_before: Option<Value>,
164    /// System state after the call
165    #[serde(default)]
166    pub system_state_after: Option<Value>,
167}
168
169/// Environment step event (Gymnasium/OpenAI Gym style).
170#[derive(Debug, Clone, Default, Serialize, Deserialize)]
171pub struct EnvironmentEvent {
172    /// Base event fields
173    #[serde(flatten)]
174    pub base: BaseEventFields,
175    /// Reward signal
176    #[serde(default)]
177    pub reward: f64,
178    /// Episode terminated flag
179    #[serde(default)]
180    pub terminated: bool,
181    /// Episode truncated flag
182    #[serde(default)]
183    pub truncated: bool,
184    /// System state before step
185    #[serde(default)]
186    pub system_state_before: Option<Value>,
187    /// System state after step (observations)
188    #[serde(default)]
189    pub system_state_after: Option<Value>,
190}
191
192/// Runtime/action selection event.
193#[derive(Debug, Clone, Default, Serialize, Deserialize)]
194pub struct RuntimeEvent {
195    /// Base event fields
196    #[serde(flatten)]
197    pub base: BaseEventFields,
198    /// Action indices/selections
199    #[serde(default)]
200    pub actions: Vec<i32>,
201}
202
203/// Unified event type using tagged enum.
204#[derive(Debug, Clone, Serialize, Deserialize)]
205#[serde(tag = "event_type", rename_all = "snake_case")]
206pub enum TracingEvent {
207    /// LLM/CAIS call
208    Cais(LMCAISEvent),
209    /// Environment step
210    Environment(EnvironmentEvent),
211    /// Runtime action
212    Runtime(RuntimeEvent),
213}
214
215impl TracingEvent {
216    /// Get the event type.
217    pub fn event_type(&self) -> EventType {
218        match self {
219            TracingEvent::Cais(_) => EventType::Cais,
220            TracingEvent::Environment(_) => EventType::Environment,
221            TracingEvent::Runtime(_) => EventType::Runtime,
222        }
223    }
224
225    /// Get the base event fields.
226    pub fn base(&self) -> &BaseEventFields {
227        match self {
228            TracingEvent::Cais(e) => &e.base,
229            TracingEvent::Environment(e) => &e.base,
230            TracingEvent::Runtime(e) => &e.base,
231        }
232    }
233
234    /// Get the time record.
235    pub fn time_record(&self) -> &TimeRecord {
236        &self.base().time_record
237    }
238
239    /// Get the system instance ID.
240    pub fn system_instance_id(&self) -> &str {
241        &self.base().system_instance_id
242    }
243}
244
245// ============================================================================
246// LLM CALL RECORDS
247// ============================================================================
248
249/// Token usage statistics.
250#[derive(Debug, Clone, Default, Serialize, Deserialize)]
251pub struct LLMUsage {
252    #[serde(default)]
253    pub input_tokens: Option<i32>,
254    #[serde(default)]
255    pub output_tokens: Option<i32>,
256    #[serde(default)]
257    pub total_tokens: Option<i32>,
258    #[serde(default)]
259    pub reasoning_tokens: Option<i32>,
260    #[serde(default)]
261    pub reasoning_input_tokens: Option<i32>,
262    #[serde(default)]
263    pub reasoning_output_tokens: Option<i32>,
264    #[serde(default)]
265    pub cache_read_tokens: Option<i32>,
266    #[serde(default)]
267    pub cache_write_tokens: Option<i32>,
268    #[serde(default)]
269    pub billable_input_tokens: Option<i32>,
270    #[serde(default)]
271    pub billable_output_tokens: Option<i32>,
272    #[serde(default)]
273    pub cost_usd: Option<f64>,
274}
275
276/// Provider request parameters.
277#[derive(Debug, Clone, Default, Serialize, Deserialize)]
278pub struct LLMRequestParams {
279    #[serde(default)]
280    pub temperature: Option<f64>,
281    #[serde(default)]
282    pub top_p: Option<f64>,
283    #[serde(default)]
284    pub max_tokens: Option<i32>,
285    #[serde(default)]
286    pub stop: Option<Vec<String>>,
287    #[serde(default)]
288    pub top_k: Option<i32>,
289    #[serde(default)]
290    pub presence_penalty: Option<f64>,
291    #[serde(default)]
292    pub frequency_penalty: Option<f64>,
293    #[serde(default)]
294    pub repetition_penalty: Option<f64>,
295    #[serde(default)]
296    pub seed: Option<i32>,
297    #[serde(default)]
298    pub n: Option<i32>,
299    #[serde(default)]
300    pub best_of: Option<i32>,
301    #[serde(default)]
302    pub response_format: Option<Value>,
303    #[serde(default)]
304    pub json_mode: Option<bool>,
305    #[serde(default)]
306    pub tool_config: Option<Value>,
307    #[serde(default)]
308    pub raw_params: HashMap<String, Value>,
309}
310
311/// LLM message content part.
312#[derive(Debug, Clone, Default, Serialize, Deserialize)]
313pub struct LLMContentPart {
314    /// Content type (text, image, audio, etc.)
315    #[serde(rename = "type", default)]
316    pub content_type: String,
317    /// Text content
318    #[serde(default)]
319    pub text: Option<String>,
320    /// Generic data payload
321    #[serde(default)]
322    pub data: Option<Value>,
323    /// MIME type for media
324    #[serde(default)]
325    pub mime_type: Option<String>,
326    #[serde(default)]
327    pub uri: Option<String>,
328    #[serde(default)]
329    pub base64_data: Option<String>,
330    #[serde(default)]
331    pub size_bytes: Option<i64>,
332    #[serde(default)]
333    pub sha256: Option<String>,
334    #[serde(default)]
335    pub width: Option<i32>,
336    #[serde(default)]
337    pub height: Option<i32>,
338    #[serde(default)]
339    pub duration_ms: Option<i32>,
340    #[serde(default)]
341    pub sample_rate: Option<i32>,
342    #[serde(default)]
343    pub channels: Option<i32>,
344    #[serde(default)]
345    pub language: Option<String>,
346}
347
348impl LLMContentPart {
349    /// Create a text content part.
350    pub fn text(text: impl Into<String>) -> Self {
351        Self {
352            content_type: "text".to_string(),
353            text: Some(text.into()),
354            data: None,
355            mime_type: None,
356            uri: None,
357            base64_data: None,
358            size_bytes: None,
359            sha256: None,
360            width: None,
361            height: None,
362            duration_ms: None,
363            sample_rate: None,
364            channels: None,
365            language: None,
366        }
367    }
368}
369
370/// LLM message.
371#[derive(Debug, Clone, Default, Serialize, Deserialize)]
372pub struct LLMMessage {
373    /// Role (system, user, assistant, tool)
374    #[serde(default)]
375    pub role: String,
376    /// Message content parts
377    #[serde(default)]
378    pub parts: Vec<LLMContentPart>,
379    /// Optional message name
380    #[serde(default)]
381    pub name: Option<String>,
382    /// Tool call ID for tool messages
383    #[serde(default)]
384    pub tool_call_id: Option<String>,
385    /// Additional metadata
386    #[serde(default)]
387    pub metadata: HashMap<String, Value>,
388}
389
390impl LLMMessage {
391    /// Create a simple text message.
392    pub fn new(role: impl Into<String>, text: impl Into<String>) -> Self {
393        Self {
394            role: role.into(),
395            parts: vec![LLMContentPart::text(text)],
396            name: None,
397            tool_call_id: None,
398            metadata: HashMap::new(),
399        }
400    }
401
402    /// Get the text content of the message.
403    pub fn text(&self) -> Option<&str> {
404        self.parts.iter().find_map(|p| p.text.as_deref())
405    }
406}
407
408/// Tool call specification.
409#[derive(Debug, Clone, Default, Serialize, Deserialize)]
410pub struct ToolCallSpec {
411    /// Tool/function name
412    #[serde(default)]
413    pub name: String,
414    /// Arguments as JSON string
415    #[serde(default)]
416    pub arguments_json: String,
417    /// Parsed arguments (optional)
418    #[serde(default)]
419    pub arguments: Option<Value>,
420    /// Call ID
421    #[serde(default)]
422    pub call_id: Option<String>,
423    /// Index in batch
424    #[serde(default)]
425    pub index: Option<i32>,
426    /// Parent call ID
427    #[serde(default)]
428    pub parent_call_id: Option<String>,
429    /// Additional metadata
430    #[serde(default)]
431    pub metadata: HashMap<String, Value>,
432}
433
434/// Tool call result.
435#[derive(Debug, Clone, Default, Serialize, Deserialize)]
436pub struct ToolCallResult {
437    /// Correlates to ToolCallSpec
438    #[serde(default)]
439    pub call_id: Option<String>,
440    /// Execution result text
441    #[serde(default)]
442    pub output_text: Option<String>,
443    /// Exit code
444    #[serde(default)]
445    pub exit_code: Option<i32>,
446    /// Status (ok, error)
447    #[serde(default)]
448    pub status: Option<String>,
449    /// Error message
450    #[serde(default)]
451    pub error_message: Option<String>,
452    /// Start timestamp
453    #[serde(default)]
454    pub started_at: Option<DateTime<Utc>>,
455    /// Completion timestamp
456    #[serde(default)]
457    pub completed_at: Option<DateTime<Utc>>,
458    /// Duration in milliseconds
459    #[serde(default)]
460    pub duration_ms: Option<i32>,
461    /// Additional metadata
462    #[serde(default)]
463    pub metadata: HashMap<String, Value>,
464}
465
466/// Optional streaming chunk representation.
467#[derive(Debug, Clone, Default, Serialize, Deserialize)]
468pub struct LLMChunk {
469    #[serde(default)]
470    pub sequence_index: i32,
471    #[serde(default = "Utc::now")]
472    pub received_at: DateTime<Utc>,
473    #[serde(default)]
474    pub event_type: Option<String>,
475    #[serde(default)]
476    pub choice_index: Option<i32>,
477    #[serde(default)]
478    pub raw_json: Option<String>,
479    #[serde(default)]
480    pub delta_text: Option<String>,
481    #[serde(default)]
482    pub delta: Option<Value>,
483    #[serde(default)]
484    pub metadata: HashMap<String, Value>,
485}
486
487/// Normalized LLM call record.
488#[derive(Debug, Clone, Default, Serialize, Deserialize)]
489pub struct LLMCallRecord {
490    /// Unique call ID
491    #[serde(default)]
492    pub call_id: String,
493    /// API type (chat_completions, completions, responses)
494    #[serde(default)]
495    pub api_type: String,
496    /// Provider (openai, anthropic, etc.)
497    #[serde(default)]
498    pub provider: Option<String>,
499    /// Model name
500    #[serde(default)]
501    pub model_name: String,
502    /// Schema version
503    #[serde(default)]
504    pub schema_version: Option<String>,
505    /// Call start time
506    #[serde(default)]
507    pub started_at: Option<DateTime<Utc>>,
508    /// Call completion time
509    #[serde(default)]
510    pub completed_at: Option<DateTime<Utc>>,
511    /// Latency in milliseconds
512    #[serde(default)]
513    pub latency_ms: Option<i32>,
514    /// Provider request parameters
515    #[serde(default)]
516    pub request_params: LLMRequestParams,
517    /// Input messages
518    #[serde(default)]
519    pub input_messages: Vec<LLMMessage>,
520    /// Input text (completions-style)
521    #[serde(default)]
522    pub input_text: Option<String>,
523    /// Tool choice
524    #[serde(default)]
525    pub tool_choice: Option<String>,
526    /// Output messages
527    #[serde(default)]
528    pub output_messages: Vec<LLMMessage>,
529    /// Output choices (n>1)
530    #[serde(default)]
531    pub outputs: Vec<LLMMessage>,
532    /// Output text (completions-style)
533    #[serde(default)]
534    pub output_text: Option<String>,
535    /// Tool calls in response
536    #[serde(default)]
537    pub output_tool_calls: Vec<ToolCallSpec>,
538    /// Tool execution results
539    #[serde(default)]
540    pub tool_results: Vec<ToolCallResult>,
541    /// Token usage
542    #[serde(default)]
543    pub usage: Option<LLMUsage>,
544    /// Finish reason
545    #[serde(default)]
546    pub finish_reason: Option<String>,
547    /// Choice index
548    #[serde(default)]
549    pub choice_index: Option<i32>,
550    /// Streaming chunks
551    #[serde(default)]
552    pub chunks: Option<Vec<LLMChunk>>,
553    /// Raw request JSON
554    #[serde(default)]
555    pub request_raw_json: Option<String>,
556    /// Raw response JSON
557    #[serde(default)]
558    pub response_raw_json: Option<String>,
559    /// Additional metadata
560    #[serde(default)]
561    pub metadata: HashMap<String, Value>,
562    /// Provider request ID
563    #[serde(default)]
564    pub provider_request_id: Option<String>,
565    /// Request server timing info
566    #[serde(default)]
567    pub request_server_timing: Option<Value>,
568    /// Outcome status
569    #[serde(default)]
570    pub outcome: Option<String>,
571    /// Error details
572    #[serde(default)]
573    pub error: Option<Value>,
574    /// Token trace info
575    #[serde(default)]
576    pub token_traces: Option<Vec<Value>>,
577    /// Safety metadata
578    #[serde(default)]
579    pub safety: Option<Value>,
580    /// Refusal metadata
581    #[serde(default)]
582    pub refusal: Option<Value>,
583    /// Redactions
584    #[serde(default)]
585    pub redactions: Option<Vec<Value>>,
586}
587
588// ============================================================================
589// SESSION STRUCTURE
590// ============================================================================
591
592/// Inter-system message (Markov blanket).
593#[derive(Debug, Clone, Default, Serialize, Deserialize)]
594pub struct MarkovBlanketMessage {
595    /// Message content
596    #[serde(default)]
597    pub content: MessageContent,
598    /// Message type (user, assistant, system, tool_use, tool_result)
599    #[serde(default)]
600    pub message_type: String,
601    /// Time record
602    #[serde(default)]
603    pub time_record: TimeRecord,
604    /// Additional metadata
605    #[serde(default)]
606    pub metadata: HashMap<String, Value>,
607}
608
609/// A timestep within a session.
610#[derive(Debug, Clone, Serialize, Deserialize)]
611pub struct SessionTimeStep {
612    /// Unique step ID
613    #[serde(default)]
614    pub step_id: String,
615    /// Sequential step index
616    #[serde(default)]
617    pub step_index: i32,
618    /// Step start time
619    #[serde(default = "Utc::now")]
620    pub timestamp: DateTime<Utc>,
621    /// Conversation turn number
622    #[serde(default)]
623    pub turn_number: Option<i32>,
624    /// Events in this step
625    #[serde(default)]
626    pub events: Vec<TracingEvent>,
627    /// Messages in this step
628    #[serde(default)]
629    pub markov_blanket_messages: Vec<MarkovBlanketMessage>,
630    /// Step-specific metadata
631    #[serde(default)]
632    pub step_metadata: HashMap<String, Value>,
633    /// Step completion time
634    #[serde(default)]
635    pub completed_at: Option<DateTime<Utc>>,
636}
637
638impl SessionTimeStep {
639    /// Create a new timestep.
640    pub fn new(step_id: impl Into<String>, step_index: i32) -> Self {
641        Self {
642            step_id: step_id.into(),
643            step_index,
644            timestamp: Utc::now(),
645            turn_number: None,
646            events: Vec::new(),
647            markov_blanket_messages: Vec::new(),
648            step_metadata: HashMap::new(),
649            completed_at: None,
650        }
651    }
652
653    /// Mark the timestep as complete.
654    pub fn complete(&mut self) {
655        self.completed_at = Some(Utc::now());
656    }
657}
658
659/// A complete session trace.
660#[derive(Debug, Clone, Serialize, Deserialize)]
661pub struct SessionTrace {
662    /// Session ID
663    #[serde(default)]
664    pub session_id: String,
665    /// Session creation time
666    #[serde(default = "Utc::now")]
667    pub created_at: DateTime<Utc>,
668    /// Ordered timesteps
669    #[serde(default)]
670    pub session_time_steps: Vec<SessionTimeStep>,
671    /// Flattened event history
672    #[serde(default)]
673    pub event_history: Vec<TracingEvent>,
674    /// Flattened message history
675    #[serde(default)]
676    pub markov_blanket_message_history: Vec<MarkovBlanketMessage>,
677    /// Session-level metadata
678    #[serde(default)]
679    pub metadata: HashMap<String, Value>,
680}
681
682impl SessionTrace {
683    /// Create a new session trace.
684    pub fn new(session_id: impl Into<String>) -> Self {
685        Self {
686            session_id: session_id.into(),
687            created_at: Utc::now(),
688            session_time_steps: Vec::new(),
689            event_history: Vec::new(),
690            markov_blanket_message_history: Vec::new(),
691            metadata: HashMap::new(),
692        }
693    }
694
695    /// Get the number of timesteps.
696    pub fn num_timesteps(&self) -> usize {
697        self.session_time_steps.len()
698    }
699
700    /// Get the number of events.
701    pub fn num_events(&self) -> usize {
702        self.event_history.len()
703    }
704
705    /// Get the number of messages.
706    pub fn num_messages(&self) -> usize {
707        self.markov_blanket_message_history.len()
708    }
709}
710
711// ============================================================================
712// REWARDS
713// ============================================================================
714
715/// Session-level outcome reward.
716#[derive(Debug, Clone, Default, Serialize, Deserialize)]
717pub struct OutcomeReward {
718    /// Objective key (default: "reward")
719    #[serde(default = "default_objective_key")]
720    pub objective_key: String,
721    /// Total reward value
722    pub total_reward: f64,
723    /// Number of achievements
724    #[serde(default)]
725    pub achievements_count: i32,
726    /// Total steps in session
727    #[serde(default)]
728    pub total_steps: i32,
729    /// Additional metadata
730    #[serde(default)]
731    pub reward_metadata: HashMap<String, Value>,
732    /// Annotation
733    #[serde(default)]
734    pub annotation: Option<Value>,
735}
736
737fn default_objective_key() -> String {
738    "reward".to_string()
739}
740
741/// Event-level reward.
742#[derive(Debug, Clone, Default, Serialize, Deserialize)]
743pub struct EventReward {
744    /// Objective key (default: "reward")
745    #[serde(default = "default_objective_key")]
746    pub objective_key: String,
747    /// Reward value
748    pub reward_value: f64,
749    /// Reward type (shaped, sparse, achievement, penalty, evaluator, human)
750    #[serde(default)]
751    pub reward_type: Option<String>,
752    /// Key (e.g., achievement name)
753    #[serde(default)]
754    pub key: Option<String>,
755    /// Annotation
756    #[serde(default)]
757    pub annotation: Option<Value>,
758    /// Source (environment, runner, evaluator, human)
759    #[serde(default)]
760    pub source: Option<String>,
761}
762
763#[cfg(test)]
764mod tests {
765    use super::*;
766
767    #[test]
768    fn test_time_record() {
769        let tr = TimeRecord::now();
770        assert!(tr.event_time > 0.0);
771        assert!(tr.message_time.is_none());
772    }
773
774    #[test]
775    fn test_message_content() {
776        let mc = MessageContent::from_text("hello");
777        assert_eq!(mc.as_text(), Some("hello"));
778
779        let mc = MessageContent::from_json(&serde_json::json!({"key": "value"}));
780        assert!(mc.json_payload.is_some());
781    }
782
783    #[test]
784    fn test_event_serialization() {
785        let event = TracingEvent::Cais(LMCAISEvent {
786            base: BaseEventFields::new("test-system"),
787            model_name: "gpt-4".to_string(),
788            provider: Some("openai".to_string()),
789            input_tokens: Some(100),
790            output_tokens: Some(50),
791            ..Default::default()
792        });
793
794        let json = serde_json::to_string(&event).unwrap();
795        assert!(json.contains("cais"));
796        assert!(json.contains("gpt-4"));
797
798        let parsed: TracingEvent = serde_json::from_str(&json).unwrap();
799        assert_eq!(parsed.event_type(), EventType::Cais);
800    }
801
802    #[test]
803    fn test_session_trace() {
804        let mut trace = SessionTrace::new("test-session");
805        assert_eq!(trace.num_timesteps(), 0);
806
807        let step = SessionTimeStep::new("step-1", 0);
808        trace.session_time_steps.push(step);
809        assert_eq!(trace.num_timesteps(), 1);
810    }
811
812    #[test]
813    fn test_llm_message() {
814        let msg = LLMMessage::new("user", "Hello, world!");
815        assert_eq!(msg.role, "user");
816        assert_eq!(msg.text(), Some("Hello, world!"));
817    }
818}