use crate::value::VmDictExt;
use std::collections::BTreeMap;
use crate::value::{VmError, VmValue};
use crate::vm::Vm;
use harn_builtin_macros::harn_builtin;
pub(crate) fn register_transcript_builtins(vm: &mut Vm) {
harn_vm_core::builtin_support::register_module_builtins(vm, module_path!());
}
#[harn_builtin(
exposure = "pure",
effects = [],
sig = "transcript_stats(transcript: list | dict | Transcript | SessionSnapshot | nil) -> dict",
category = "agent"
)]
fn transcript_stats_builtin(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
let transcript = args.first().cloned().unwrap_or(VmValue::Nil);
Ok(stats_to_vm(&compute_transcript_stats(&transcript)))
}
#[harn_builtin(
exposure = "pure",
effects = [],
sig = "transcript_events_by_kind(transcript: list | dict | Transcript | SessionSnapshot | nil, kind: string) -> list",
category = "agent"
)]
fn transcript_events_by_kind_builtin(
args: &[VmValue],
_out: &mut String,
) -> Result<VmValue, VmError> {
let transcript = args.first().cloned().unwrap_or(VmValue::Nil);
let kind = args.get(1).map(VmValue::display).unwrap_or_default();
let filtered = events_list(&transcript)
.into_iter()
.filter(|event| match event {
VmValue::Dict(dict) => dict
.get("kind")
.and_then(|value| match value {
VmValue::String(value) => Some(value.as_str()),
_ => None,
})
.is_some_and(|value| value == kind),
_ => false,
})
.collect();
Ok(VmValue::list(filtered))
}
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub(crate) struct TranscriptStats {
pub message_count: i64,
pub user_facing_message_count: i64,
pub tool_result_message_count: i64,
pub tool_call_count: i64,
pub assistant_message_count: i64,
pub user_message_count: i64,
pub visible_event_count: i64,
}
pub(crate) fn compute_transcript_stats(transcript: &VmValue) -> TranscriptStats {
let Some(dict) = transcript.as_dict() else {
return TranscriptStats::default();
};
let messages = list_field(dict, "messages");
let events = list_field(dict, "events");
let mut stats = TranscriptStats::default();
for message in &messages {
let Some(msg) = message.as_dict() else {
continue;
};
stats.message_count += 1;
let role = string_field(msg, "role").unwrap_or_default();
match role.as_str() {
"assistant" => {
stats.assistant_message_count += 1;
stats.user_facing_message_count += 1;
}
"user" => {
stats.user_message_count += 1;
stats.user_facing_message_count += 1;
}
"tool" | "tool_result" => {
stats.tool_result_message_count += 1;
}
_ => {}
}
stats.tool_call_count += count_tool_calls(msg);
}
for event in &events {
let Some(ev) = event.as_dict() else { continue };
if string_field(ev, "visibility").as_deref() == Some("visible") {
stats.visible_event_count += 1;
}
}
stats
}
fn events_list(transcript: &VmValue) -> Vec<VmValue> {
transcript
.as_dict()
.map(|dict| list_field(dict, "events"))
.unwrap_or_default()
}
fn list_field(dict: &crate::value::DictMap, key: &str) -> Vec<VmValue> {
match dict.get(key) {
Some(VmValue::List(list)) => (**list).clone(),
_ => Vec::new(),
}
}
fn string_field(dict: &crate::value::DictMap, key: &str) -> Option<String> {
match dict.get(key) {
Some(VmValue::String(s)) => Some(s.to_string()),
_ => None,
}
}
fn count_tool_calls(msg: &crate::value::DictMap) -> i64 {
match msg.get("tool_calls") {
Some(VmValue::List(list)) => list.len() as i64,
_ => 0,
}
}
fn stats_to_vm(stats: &TranscriptStats) -> VmValue {
let mut dict = BTreeMap::new();
dict.put_int("message_count", stats.message_count);
dict.put_int("user_facing_message_count", stats.user_facing_message_count);
dict.put_int("tool_result_message_count", stats.tool_result_message_count);
dict.put_int("tool_call_count", stats.tool_call_count);
dict.put_int("assistant_message_count", stats.assistant_message_count);
dict.put_int("user_message_count", stats.user_message_count);
dict.put_int("visible_event_count", stats.visible_event_count);
VmValue::dict(dict)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::schema::json_to_vm_value;
fn transcript_from_json(value: serde_json::Value) -> VmValue {
json_to_vm_value(&value)
}
#[test]
fn stats_nil_transcript_returns_zeroes() {
let stats = compute_transcript_stats(&VmValue::Nil);
assert_eq!(stats, TranscriptStats::default());
}
#[test]
fn stats_counts_roles_across_messages() {
let transcript = transcript_from_json(serde_json::json!({
"messages": [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "hey"},
{"role": "assistant", "content": "", "tool_calls": [{"id": "a"}, {"id": "b"}]},
{"role": "tool", "content": "ok", "tool_use_id": "a"},
{"role": "tool_result", "content": "ok", "tool_use_id": "b"},
{"role": "system", "content": "ignore me"},
],
}));
let stats = compute_transcript_stats(&transcript);
assert_eq!(stats.message_count, 6);
assert_eq!(stats.user_message_count, 1);
assert_eq!(stats.assistant_message_count, 2);
assert_eq!(stats.user_facing_message_count, 3);
assert_eq!(stats.tool_result_message_count, 2);
assert_eq!(stats.tool_call_count, 2);
}
#[test]
fn stats_counts_visible_events() {
let transcript = transcript_from_json(serde_json::json!({
"events": [
{"kind": "message", "visibility": "visible"},
{"kind": "message", "visibility": "visible"},
{"kind": "tool_result", "visibility": "hidden"},
],
}));
let stats = compute_transcript_stats(&transcript);
assert_eq!(stats.visible_event_count, 2);
}
#[test]
fn stats_tolerate_missing_keys() {
let transcript = transcript_from_json(serde_json::json!({
"messages": [
{"role": "assistant"},
{"content": "no role"},
],
}));
let stats = compute_transcript_stats(&transcript);
assert_eq!(stats.message_count, 2);
assert_eq!(stats.assistant_message_count, 1);
}
}