agtrace_runtime/ops/
analyze.rs

1use crate::Result;
2use crate::storage::{LoadOptions, SessionRepository};
3use agtrace_engine::assemble_session;
4use agtrace_index::Database;
5use agtrace_providers::create_adapter;
6use agtrace_types::EventPayload;
7use std::collections::{BTreeMap, HashMap};
8
9#[derive(Debug, Clone)]
10pub struct CorpusStats {
11    pub sample_size: usize,
12    pub total_tool_calls: usize,
13    pub total_failures: usize,
14    pub max_duration_ms: i64,
15}
16
17pub fn get_corpus_overview(
18    db: &Database,
19    project_hash: Option<&agtrace_types::ProjectHash>,
20    limit: usize,
21) -> Result<CorpusStats> {
22    let raw_sessions = db.list_sessions(project_hash, None, Some(limit))?;
23
24    let loader = SessionRepository::new(db);
25    let options = LoadOptions::default();
26
27    let mut total_tool_calls = 0;
28    let mut total_failures = 0;
29    let mut max_duration = 0i64;
30
31    for session in &raw_sessions {
32        if let Ok(events) = loader.load_events(&session.id, &options)
33            && let Some(agent_session) = assemble_session(&events)
34        {
35            for turn in &agent_session.turns {
36                for step in &turn.steps {
37                    total_tool_calls += step.tools.len();
38                    for tool_exec in &step.tools {
39                        if tool_exec.is_error {
40                            total_failures += 1;
41                        }
42                    }
43                }
44                if turn.stats.duration_ms > max_duration {
45                    max_duration = turn.stats.duration_ms;
46                }
47            }
48        }
49    }
50
51    Ok(CorpusStats {
52        sample_size: raw_sessions.len(),
53        total_tool_calls,
54        total_failures,
55        max_duration_ms: max_duration,
56    })
57}
58
59#[derive(Debug, Clone)]
60pub struct ToolSample {
61    pub arguments: String,
62    pub result: Option<String>,
63}
64
65#[derive(Debug, Clone)]
66pub struct ToolInfo {
67    pub tool_name: String,
68    pub origin: Option<String>,
69    pub kind: Option<String>,
70}
71
72pub type ProviderStats =
73    BTreeMap<String, (BTreeMap<String, (usize, Option<ToolSample>)>, Vec<ToolInfo>)>;
74
75pub struct StatsResult {
76    pub total_sessions: usize,
77    pub provider_stats: ProviderStats,
78}
79
80pub fn collect_tool_stats(
81    db: &Database,
82    limit: Option<usize>,
83    provider: Option<String>,
84) -> Result<StatsResult> {
85    let sessions = db.list_sessions(None, provider.as_deref(), limit)?;
86    let total_sessions = sessions.len();
87
88    let loader = SessionRepository::new(db);
89    let options = LoadOptions::default();
90
91    let mut stats: HashMap<String, HashMap<String, (usize, Option<ToolSample>)>> = HashMap::new();
92
93    for session in &sessions {
94        let events = match loader.load_events(&session.id, &options) {
95            Ok(events) => events,
96            Err(_) => continue,
97        };
98
99        let mut tool_results = HashMap::new();
100        for event in &events {
101            if let EventPayload::ToolResult(result) = &event.payload {
102                tool_results.insert(result.tool_call_id, result.output.clone());
103            }
104        }
105
106        let provider = &session.provider;
107        for event in events {
108            if let EventPayload::ToolCall(tool_call) = &event.payload {
109                let provider_stats = stats.entry(provider.clone()).or_default();
110                let tool_entry = provider_stats
111                    .entry(tool_call.name().to_string())
112                    .or_insert((0, None));
113
114                tool_entry.0 += 1;
115
116                if tool_entry.1.is_none() {
117                    let result = tool_results.get(&event.id).cloned();
118                    // Serialize entire tool_call to extract arguments field
119                    let arguments = serde_json::to_value(tool_call)
120                        .ok()
121                        .and_then(|v| v.get("arguments").cloned())
122                        .and_then(|v| serde_json::to_string(&v).ok())
123                        .unwrap_or_else(|| String::from("(failed to serialize)"));
124                    tool_entry.1 = Some(ToolSample { arguments, result });
125                }
126            }
127        }
128    }
129
130    let provider_stats: ProviderStats = stats
131        .into_iter()
132        .map(|(provider_name, tools)| {
133            let sorted_tools: BTreeMap<_, _> = tools.into_iter().collect();
134
135            let classifications: Vec<ToolInfo> = sorted_tools
136                .keys()
137                .map(|tool_name| {
138                    let (origin, kind) = if let Ok(adapter) = create_adapter(&provider_name) {
139                        let (o, k) = adapter.mapper.classify(tool_name);
140                        (Some(format!("{:?}", o)), Some(format!("{:?}", k)))
141                    } else {
142                        (None, None)
143                    };
144
145                    ToolInfo {
146                        tool_name: tool_name.clone(),
147                        origin,
148                        kind,
149                    }
150                })
151                .collect();
152
153            (provider_name, (sorted_tools, classifications))
154        })
155        .collect();
156
157    Ok(StatsResult {
158        total_sessions,
159        provider_stats,
160    })
161}
agtrace_runtime/ops/analyze.rs

agtrace_runtime/ops/
analyze.rs
