Skip to main content

obol_core/transcript/
kimi.rs

1//! Kimi Code wire.jsonl -> Vec<MessageUsage>.
2//! Targets quorum's `usage.record` rows (full input/cache fidelity + model), not
3//! agentsview's lower-fidelity StatusUpdate path. Prefer `usageScope:"turn"` rows;
4//! fall back to the latest `session` row. Never mix turn + session (double-counts).
5
6use crate::error::ObolError;
7use crate::model::{MessageUsage, Provider};
8use serde_json::Value;
9
10pub fn parse(bytes: &[u8]) -> Result<Vec<MessageUsage>, ObolError> {
11    let text = std::str::from_utf8(bytes).map_err(|e| ObolError::MalformedTranscript {
12        line: 0,
13        msg: e.to_string(),
14    })?;
15    let mut turns: Vec<Value> = Vec::new();
16    let mut sessions: Vec<Value> = Vec::new();
17    for line in text.lines() {
18        let line = line.trim();
19        if line.is_empty() {
20            continue;
21        }
22        let v: Value = match serde_json::from_str(line) {
23            Ok(v) => v,
24            Err(_) => continue,
25        };
26        if v.get("type").and_then(Value::as_str) != Some("usage.record") {
27            continue;
28        }
29        match v.get("usageScope").and_then(Value::as_str) {
30            Some("turn") => turns.push(v),
31            Some("session") => sessions.push(v),
32            _ => {}
33        }
34    }
35    let selected: Vec<Value> = if !turns.is_empty() {
36        turns
37    } else {
38        match sessions
39            .into_iter()
40            .max_by_key(|r| r.get("time").and_then(Value::as_i64).unwrap_or(i64::MIN))
41        {
42            Some(latest) => vec![latest],
43            None => Vec::new(),
44        }
45    };
46    let mut out = Vec::new();
47    for row in &selected {
48        let usage = match row.get("usage") {
49            Some(u) if u.is_object() => u,
50            _ => continue,
51        };
52        let g = |k: &str| usage.get(k).and_then(Value::as_u64).unwrap_or(0);
53        let input = g("inputOther");
54        let cache_read = g("inputCacheRead");
55        let cache_create = g("inputCacheCreation");
56        out.push(MessageUsage {
57            model: row
58                .get("model")
59                .and_then(Value::as_str)
60                .unwrap_or("")
61                .to_string(),
62            provider: Provider::Other("moonshot".into()),
63            namespace: "litellm".into(),
64            input_uncached: input,
65            cache_read,
66            cache_write_5m: cache_create,
67            cache_write_1h: 0,
68            output: g("output"),
69            request_input_tokens: input + cache_read + cache_create,
70            service_tier: None,
71        });
72    }
73    Ok(out)
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    #[test]
81    fn prefers_turn_rows_over_session() {
82        let u = parse(include_bytes!("../../tests/fixtures/kimi-mini.jsonl")).unwrap();
83        assert_eq!(
84            u.len(),
85            2,
86            "session row must be ignored when turns exist: {u:?}"
87        );
88        assert_eq!(u[0].model, "kimi-for-coding");
89        assert_eq!(u[0].input_uncached, 10);
90        assert_eq!(u[0].cache_read, 20);
91        assert_eq!(u[0].cache_write_5m, 30);
92        assert_eq!(u[0].output, 40);
93        assert_eq!(u[0].request_input_tokens, 60);
94    }
95}