mempal_runtime/ingest/
normalize.rs

1use serde_json::Value;
2use thiserror::Error;
3
4use super::detect::{Format, extract_content_text, extract_message_text};
5use super::noise::{strip_claude_jsonl_noise, strip_codex_rollout_noise};
6
7pub const CURRENT_NORMALIZE_VERSION: u32 = 2;
8
9pub type Result<T> = std::result::Result<T, NormalizeError>;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub struct NormalizeOptions {
13    pub strip_noise: bool,
14}
15
16impl Default for NormalizeOptions {
17    fn default() -> Self {
18        Self { strip_noise: true }
19    }
20}
21
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct NormalizeOutput {
24    pub content: String,
25    pub noise_bytes_stripped: Option<u64>,
26}
27
28#[derive(Debug, Error)]
29pub enum NormalizeError {
30    #[error(transparent)]
31    Json(#[from] serde_json::Error),
32    #[error("unsupported ChatGPT JSON shape")]
33    UnsupportedChatGptShape,
34}
35
36pub fn normalize_content(content: &str, format: Format) -> Result<String> {
37    Ok(normalize_content_with_options(content, format, NormalizeOptions::default())?.content)
38}
39
40pub fn normalize_content_with_options(
41    content: &str,
42    format: Format,
43    options: NormalizeOptions,
44) -> Result<NormalizeOutput> {
45    match format {
46        Format::PlainText => Ok(NormalizeOutput {
47            content: content.trim().to_string(),
48            noise_bytes_stripped: None,
49        }),
50        Format::ClaudeJsonl => normalize_claude_jsonl(content, options.strip_noise),
51        Format::ChatGptJson => Ok(NormalizeOutput {
52            content: normalize_chatgpt_json(content)?,
53            noise_bytes_stripped: None,
54        }),
55        Format::CodexJsonl => normalize_codex_jsonl(content, options.strip_noise),
56        Format::SlackJson => Ok(NormalizeOutput {
57            content: normalize_slack_json(content)?,
58            noise_bytes_stripped: None,
59        }),
60    }
61}
62
63fn normalize_claude_jsonl(content: &str, strip_noise: bool) -> Result<NormalizeOutput> {
64    let mut lines = Vec::new();
65    let mut noise_bytes_stripped = 0_u64;
66
67    for raw_line in content
68        .lines()
69        .map(str::trim)
70        .filter(|line| !line.is_empty())
71    {
72        let value: Value = serde_json::from_str(raw_line)?;
73        let role = value
74            .get("type")
75            .and_then(Value::as_str)
76            .unwrap_or("assistant");
77        let message = extract_message_text(&value).unwrap_or_default();
78        let message = message.trim();
79        let message = if strip_noise {
80            let stripped = strip_claude_jsonl_noise(message);
81            noise_bytes_stripped += message.len().saturating_sub(stripped.len()) as u64;
82            stripped
83        } else {
84            message.to_string()
85        };
86
87        if message.trim().is_empty() {
88            continue;
89        }
90
91        if matches!(role, "human" | "user") {
92            lines.push(format!("> {}", message.trim()));
93        } else {
94            lines.push(message.trim().to_string());
95        }
96    }
97
98    Ok(NormalizeOutput {
99        content: lines.join("\n"),
100        noise_bytes_stripped: strip_noise.then_some(noise_bytes_stripped),
101    })
102}
103
104fn normalize_chatgpt_json(content: &str) -> Result<String> {
105    let value: Value = serde_json::from_str(content)?;
106
107    if let Some(messages) = value.as_array() {
108        return normalize_chatgpt_messages(messages);
109    }
110
111    if let Some(messages) = value.get("messages").and_then(Value::as_array) {
112        return normalize_chatgpt_messages(messages);
113    }
114
115    if let Some(mapping) = value.get("mapping").and_then(Value::as_object) {
116        let mut ordered = Vec::new();
117        if let Some(root_id) = find_root_node(mapping) {
118            collect_messages_dfs(mapping, &root_id, &mut ordered);
119        }
120
121        return Ok(render_transcript(ordered));
122    }
123
124    Err(NormalizeError::UnsupportedChatGptShape)
125}
126
127fn normalize_chatgpt_messages(messages: &[Value]) -> Result<String> {
128    let transcript = render_transcript(messages.iter().filter_map(|message| {
129        let role = message.get("role").and_then(Value::as_str)?;
130        let content = message.get("content").and_then(extract_content_text)?;
131        Some((role.to_string(), content))
132    }));
133
134    Ok(transcript)
135}
136
137fn find_root_node(mapping: &serde_json::Map<String, Value>) -> Option<String> {
138    mapping
139        .iter()
140        .find(|(_, node)| {
141            node.get("parent")
142                .is_none_or(|p| p.is_null() || p.as_str() == Some(""))
143        })
144        .map(|(id, _)| id.clone())
145}
146
147fn collect_messages_dfs(
148    mapping: &serde_json::Map<String, Value>,
149    node_id: &str,
150    result: &mut Vec<(String, String)>,
151) {
152    let Some(node) = mapping.get(node_id) else {
153        return;
154    };
155
156    if let Some(message) = node.get("message") {
157        let role = message
158            .get("author")
159            .and_then(|author| author.get("role"))
160            .and_then(Value::as_str);
161        let content = message.get("content").and_then(extract_content_text);
162        if let (Some(role), Some(content)) = (role, content) {
163            result.push((role.to_string(), content));
164        }
165    }
166
167    if let Some(children) = node.get("children").and_then(Value::as_array) {
168        for child in children {
169            if let Some(child_id) = child.as_str() {
170                collect_messages_dfs(mapping, child_id, result);
171            }
172        }
173    }
174}
175
176fn normalize_codex_jsonl(content: &str, strip_noise: bool) -> Result<NormalizeOutput> {
177    let mut pairs: Vec<(String, String)> = Vec::new();
178    let mut noise_bytes_stripped = 0_u64;
179
180    for line in content.lines().map(str::trim).filter(|l| !l.is_empty()) {
181        let value: Value = serde_json::from_str(line)?;
182        if value.get("type").and_then(Value::as_str) != Some("event_msg") {
183            continue;
184        }
185        let Some(payload) = value.get("payload") else {
186            continue;
187        };
188        let msg_type = payload.get("type").and_then(Value::as_str).unwrap_or("");
189        let message = payload
190            .get("message")
191            .and_then(Value::as_str)
192            .unwrap_or("")
193            .trim();
194        let message = if strip_noise {
195            let stripped = strip_codex_rollout_noise(message);
196            noise_bytes_stripped += message.len().saturating_sub(stripped.len()) as u64;
197            stripped
198        } else {
199            message.to_string()
200        };
201        if message.is_empty() {
202            continue;
203        }
204        match msg_type {
205            "user_message" => pairs.push(("user".to_string(), message)),
206            "agent_message" => pairs.push(("assistant".to_string(), message)),
207            _ => {}
208        }
209    }
210
211    Ok(NormalizeOutput {
212        content: render_transcript(pairs),
213        noise_bytes_stripped: strip_noise.then_some(noise_bytes_stripped),
214    })
215}
216
217fn normalize_slack_json(content: &str) -> Result<String> {
218    let value: Value = serde_json::from_str(content)?;
219    let messages = value
220        .as_array()
221        .ok_or(NormalizeError::UnsupportedChatGptShape)?;
222
223    let mut speakers: Vec<String> = Vec::new();
224    let mut pairs: Vec<(String, String)> = Vec::new();
225
226    for msg in messages {
227        if msg.get("type").and_then(Value::as_str) != Some("message") {
228            continue;
229        }
230        let speaker = msg
231            .get("user")
232            .or_else(|| msg.get("username"))
233            .and_then(Value::as_str)
234            .unwrap_or("unknown")
235            .to_string();
236        let text = msg.get("text").and_then(Value::as_str).unwrap_or("").trim();
237        if text.is_empty() {
238            continue;
239        }
240
241        // First speaker = user, second = assistant
242        if !speakers.contains(&speaker) {
243            speakers.push(speaker.clone());
244        }
245        let role = if speakers.first() == Some(&speaker) {
246            "user"
247        } else {
248            "assistant"
249        };
250        pairs.push((role.to_string(), text.to_string()));
251    }
252
253    Ok(render_transcript(pairs))
254}
255
256fn render_transcript(items: impl IntoIterator<Item = (String, String)>) -> String {
257    let mut lines = Vec::new();
258
259    for (role, content) in items {
260        if content.trim().is_empty() {
261            continue;
262        }
263
264        if matches!(role.as_str(), "user" | "human") {
265            lines.push(format!("> {}", content.trim()));
266        } else {
267            lines.push(content.trim().to_string());
268        }
269    }
270
271    lines.join("\n")
272}
mempal_runtime/ingest/normalize.rs

mempal_runtime/ingest/
normalize.rs
