Skip to main content

agx_core/
generic.rs

1use crate::timeline::{
2    Step, Usage, assistant_text_step, attach_usage_to_first, compute_durations, pretty_json,
3    tool_result_step, tool_use_step, user_text_step,
4};
5use anyhow::{Context, Result};
6use serde::Deserialize;
7use std::collections::HashMap;
8use std::path::Path;
9
10#[derive(Debug, Deserialize)]
11struct Conversation {
12    #[serde(default)]
13    messages: Vec<Message>,
14}
15
16#[derive(Debug, Deserialize)]
17struct Message {
18    #[serde(default)]
19    role: String,
20    #[serde(default)]
21    content: serde_json::Value,
22    #[serde(default)]
23    tool_calls: Vec<ToolCall>,
24    #[serde(default)]
25    tool_call_id: Option<String>,
26    #[serde(default)]
27    model: Option<String>,
28    /// OpenAI-compatible usage shape. Per-message where SDKs export it that
29    /// way; `None` otherwise.
30    #[serde(default)]
31    usage: Option<OpenAiUsage>,
32}
33
34#[derive(Debug, Deserialize)]
35struct OpenAiUsage {
36    #[serde(default)]
37    prompt_tokens: Option<u64>,
38    #[serde(default)]
39    completion_tokens: Option<u64>,
40    #[serde(default)]
41    cached_tokens: Option<u64>,
42}
43
44#[derive(Debug, Deserialize)]
45struct ToolCall {
46    #[serde(default)]
47    id: String,
48    #[serde(default)]
49    function: ToolFunction,
50}
51
52#[derive(Debug, Default, Deserialize)]
53struct ToolFunction {
54    #[serde(default)]
55    name: String,
56    #[serde(default)]
57    arguments: String,
58}
59
60pub fn load(path: &Path) -> Result<Vec<Step>> {
61    let content = std::fs::read_to_string(path)
62        .with_context(|| format!("reading conversation file: {}", path.display()))?;
63    let conv: Conversation = serde_json::from_str(&content)
64        .with_context(|| format!("parsing conversation file: {}", path.display()))?;
65
66    let tool_meta = collect_tool_meta(&conv.messages);
67    let mut steps = Vec::new();
68    for msg in &conv.messages {
69        match msg.role.as_str() {
70            "user" => {
71                let text = extract_text(&msg.content);
72                if !text.trim().is_empty() {
73                    steps.push(user_text_step(&text));
74                }
75            }
76            "assistant" => {
77                let first_idx = steps.len();
78                let text = extract_text(&msg.content);
79                if !text.trim().is_empty() {
80                    steps.push(assistant_text_step(&text));
81                }
82                for tc in &msg.tool_calls {
83                    let input_pretty = prettify_arguments(&tc.function.arguments);
84                    steps.push(tool_use_step(&tc.id, &tc.function.name, &input_pretty));
85                }
86                if steps.len() > first_idx {
87                    let usage = msg
88                        .usage
89                        .as_ref()
90                        .map(|u| Usage {
91                            tokens_in: u.prompt_tokens,
92                            tokens_out: u.completion_tokens,
93                            cache_read: u.cached_tokens,
94                            cache_create: None,
95                        })
96                        .unwrap_or_default();
97                    attach_usage_to_first(&mut steps, first_idx, msg.model.as_deref(), &usage);
98                }
99            }
100            "tool" => {
101                let result_text = extract_text(&msg.content);
102                let call_id = msg.tool_call_id.as_deref().unwrap_or("");
103                let meta = tool_meta.get(call_id);
104                steps.push(tool_result_step(
105                    call_id,
106                    &result_text,
107                    meta.map(|m| m.name.as_str()),
108                    meta.map(|m| m.input_pretty.as_str()),
109                ));
110            }
111            // System prompts and unknown roles — skip
112            _ => {}
113        }
114    }
115    compute_durations(&mut steps);
116    Ok(steps)
117}
118
119#[derive(Debug, Clone)]
120struct ToolMeta {
121    name: String,
122    input_pretty: String,
123}
124
125fn collect_tool_meta(messages: &[Message]) -> HashMap<String, ToolMeta> {
126    let mut map = HashMap::new();
127    for msg in messages {
128        if msg.role != "assistant" {
129            continue;
130        }
131        for tc in &msg.tool_calls {
132            map.insert(
133                tc.id.clone(),
134                ToolMeta {
135                    name: tc.function.name.clone(),
136                    input_pretty: prettify_arguments(&tc.function.arguments),
137                },
138            );
139        }
140    }
141    map
142}
143
144fn extract_text(content: &serde_json::Value) -> String {
145    if let Some(s) = content.as_str() {
146        return s.to_string();
147    }
148    if let Some(arr) = content.as_array() {
149        return arr
150            .iter()
151            .filter_map(|item| {
152                if item.get("type").and_then(|t| t.as_str()) == Some("text") {
153                    item.get("text").and_then(|t| t.as_str())
154                } else {
155                    None
156                }
157            })
158            .collect::<Vec<_>>()
159            .join("\n");
160    }
161    String::new()
162}
163
164fn prettify_arguments(args: &str) -> String {
165    if args.is_empty() {
166        return String::new();
167    }
168    match serde_json::from_str::<serde_json::Value>(args) {
169        Ok(v) => pretty_json(&v),
170        Err(_) => args.to_string(),
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177    use crate::timeline::StepKind;
178    use std::io::Write;
179    use tempfile::NamedTempFile;
180
181    fn write_file(content: &str) -> NamedTempFile {
182        let mut f = NamedTempFile::new().unwrap();
183        f.write_all(content.as_bytes()).unwrap();
184        f
185    }
186
187    #[test]
188    fn parses_user_and_assistant_messages() {
189        let json = r#"{"messages":[
190            {"role":"user","content":"hello"},
191            {"role":"assistant","content":"hi there"}
192        ]}"#;
193        let f = write_file(json);
194        let steps = load(f.path()).unwrap();
195        assert_eq!(steps.len(), 2);
196        assert_eq!(steps[0].kind, StepKind::UserText);
197        assert_eq!(steps[1].kind, StepKind::AssistantText);
198    }
199
200    #[test]
201    fn parses_tool_calls_and_results() {
202        let json = r#"{"messages":[
203            {"role":"assistant","content":"","tool_calls":[
204                {"id":"call_1","function":{"name":"search","arguments":"{\"q\":\"test\"}"}}
205            ]},
206            {"role":"tool","tool_call_id":"call_1","content":"found 3 results"}
207        ]}"#;
208        let f = write_file(json);
209        let steps = load(f.path()).unwrap();
210        assert_eq!(steps.len(), 2);
211        assert_eq!(steps[0].kind, StepKind::ToolUse);
212        assert!(steps[0].detail.contains("search"));
213        assert_eq!(steps[1].kind, StepKind::ToolResult);
214        assert!(steps[1].label.contains("search"));
215        assert!(steps[1].detail.contains("found 3 results"));
216    }
217
218    #[test]
219    fn skips_system_messages() {
220        let json = r#"{"messages":[
221            {"role":"system","content":"you are helpful"},
222            {"role":"user","content":"hi"}
223        ]}"#;
224        let f = write_file(json);
225        let steps = load(f.path()).unwrap();
226        assert_eq!(steps.len(), 1);
227        assert_eq!(steps[0].kind, StepKind::UserText);
228    }
229
230    #[test]
231    fn parses_openai_usage_and_model_on_assistant_message() {
232        let json = r#"{"messages":[
233            {"role":"assistant","content":"ok","model":"gpt-5","usage":{"prompt_tokens":42,"completion_tokens":17,"cached_tokens":9}}
234        ]}"#;
235        let f = write_file(json);
236        let steps = load(f.path()).unwrap();
237        assert_eq!(steps.len(), 1);
238        assert_eq!(steps[0].model.as_deref(), Some("gpt-5"));
239        assert_eq!(steps[0].tokens_in, Some(42));
240        assert_eq!(steps[0].tokens_out, Some(17));
241        assert_eq!(steps[0].cache_read, Some(9));
242    }
243
244    #[test]
245    fn user_message_with_usage_field_is_ignored() {
246        let json = r#"{"messages":[
247            {"role":"user","content":"hi","usage":{"prompt_tokens":5}}
248        ]}"#;
249        let f = write_file(json);
250        let steps = load(f.path()).unwrap();
251        assert_eq!(steps[0].tokens_in, None);
252    }
253
254    #[test]
255    fn handles_array_content_format() {
256        let json = r#"{"messages":[
257            {"role":"user","content":[{"type":"text","text":"hello world"}]}
258        ]}"#;
259        let f = write_file(json);
260        let steps = load(f.path()).unwrap();
261        assert_eq!(steps.len(), 1);
262        assert!(steps[0].detail.contains("hello world"));
263    }
264}