kaizen/interchange/
atif.rs1mod error;
5mod types;
6
7use super::jsonl::JsonlEvent;
8pub use error::AtifImportError;
9use std::collections::BTreeMap;
10pub use types::{
11 AtifDocument, AtifEvent, AtifSession, InterchangeEvent, InterchangeSession, InterchangeTrace,
12};
13
14pub const ATIF_FORMAT: &str = "atif";
15pub const ATIF_VERSION: u16 = 1;
16
17pub fn export_atif(trace: &InterchangeTrace) -> AtifDocument {
18 AtifDocument {
19 format: ATIF_FORMAT.into(),
20 version: ATIF_VERSION,
21 session: atif_session(&trace.session),
22 events: trace.events.iter().map(atif_event).collect(),
23 }
24}
25
26pub fn import_atif(doc: &AtifDocument) -> Result<InterchangeTrace, AtifImportError> {
27 validate_document(doc)?;
28 Ok(InterchangeTrace {
29 session: interchange_session(&doc.session),
30 events: doc
31 .events
32 .iter()
33 .map(|event| interchange_event(&doc.session.id, event))
34 .collect(),
35 })
36}
37
38pub fn trace_from_jsonl(session: InterchangeSession, events: Vec<JsonlEvent>) -> InterchangeTrace {
39 InterchangeTrace {
40 session,
41 events: events.into_iter().map(InterchangeEvent::from).collect(),
42 }
43}
44
45impl From<JsonlEvent> for InterchangeEvent {
46 fn from(event: JsonlEvent) -> Self {
47 Self {
48 session_id: event.session_id,
49 seq: event.seq,
50 ts_ms: event.ts_ms,
51 kind: event.kind,
52 source: event.source,
53 tool: event.tool,
54 tool_call_id: event.tool_call_id,
55 payload: event.payload,
56 attributes: BTreeMap::new(),
57 }
58 }
59}
60
61fn validate_document(doc: &AtifDocument) -> Result<(), AtifImportError> {
62 validate_header(doc)?;
63 doc.events
64 .iter()
65 .find(|event| !event.id.starts_with(&format!("{}:", doc.session.id)))
66 .map_or(Ok(()), |event| Err(mismatch(&event.id, &doc.session.id)))
67}
68
69fn validate_header(doc: &AtifDocument) -> Result<(), AtifImportError> {
70 if doc.format != ATIF_FORMAT {
71 return Err(AtifImportError::UnsupportedFormat(doc.format.clone()));
72 }
73 (doc.version == ATIF_VERSION)
74 .then_some(())
75 .ok_or(AtifImportError::UnsupportedVersion(doc.version))
76}
77
78fn atif_session(session: &InterchangeSession) -> AtifSession {
79 AtifSession {
80 id: session.id.clone(),
81 agent: session.agent.clone(),
82 model: session.model.clone(),
83 workspace: session.workspace.clone(),
84 started_at_ms: session.started_at_ms,
85 ended_at_ms: session.ended_at_ms,
86 attributes: session.attributes.clone(),
87 }
88}
89
90fn interchange_session(session: &AtifSession) -> InterchangeSession {
91 InterchangeSession {
92 id: session.id.clone(),
93 agent: session.agent.clone(),
94 model: session.model.clone(),
95 workspace: session.workspace.clone(),
96 started_at_ms: session.started_at_ms,
97 ended_at_ms: session.ended_at_ms,
98 attributes: session.attributes.clone(),
99 }
100}
101
102fn atif_event(event: &InterchangeEvent) -> AtifEvent {
103 AtifEvent {
104 id: event_id(&event.session_id, event.seq),
105 sequence: event.seq,
106 timestamp_ms: event.ts_ms,
107 event_type: event.kind.clone(),
108 source: event.source.clone(),
109 tool: event.tool.clone(),
110 tool_call_id: event.tool_call_id.clone(),
111 payload: event.payload.clone(),
112 attributes: event.attributes.clone(),
113 }
114}
115
116fn interchange_event(session_id: &str, event: &AtifEvent) -> InterchangeEvent {
117 InterchangeEvent {
118 session_id: session_id.into(),
119 seq: event.sequence,
120 ts_ms: event.timestamp_ms,
121 kind: event.event_type.clone(),
122 source: event.source.clone(),
123 tool: event.tool.clone(),
124 tool_call_id: event.tool_call_id.clone(),
125 payload: event.payload.clone(),
126 attributes: event.attributes.clone(),
127 }
128}
129
130fn event_id(session_id: &str, seq: u64) -> String {
131 format!("{session_id}:{seq}")
132}
133
134fn mismatch(event_id: &str, session_id: &str) -> AtifImportError {
135 AtifImportError::SessionMismatch {
136 event_id: event_id.into(),
137 session_id: session_id.into(),
138 }
139}