agtrace_runtime/ops/
analyze.rs

1use crate::storage::{LoadOptions, SessionRepository};
2use agtrace_engine::assemble_session;
3use agtrace_index::Database;
4use agtrace_providers::create_adapter;
5use agtrace_types::EventPayload;
6use anyhow::Result;
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, 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, limit.unwrap_or(100000))?;
86    let sessions: Vec<_> = sessions
87        .into_iter()
88        .filter(|s| {
89            if let Some(ref src) = provider {
90                &s.provider == src
91            } else {
92                true
93            }
94        })
95        .collect();
96
97    let total_sessions = sessions.len();
98
99    let loader = SessionRepository::new(db);
100    let options = LoadOptions::default();
101
102    let mut stats: HashMap<String, HashMap<String, (usize, Option<ToolSample>)>> = HashMap::new();
103
104    for session in &sessions {
105        let events = match loader.load_events(&session.id, &options) {
106            Ok(events) => events,
107            Err(_) => continue,
108        };
109
110        let mut tool_results = HashMap::new();
111        for event in &events {
112            if let EventPayload::ToolResult(result) = &event.payload {
113                tool_results.insert(result.tool_call_id, result.output.clone());
114            }
115        }
116
117        let provider = &session.provider;
118        for event in events {
119            if let EventPayload::ToolCall(tool_call) = &event.payload {
120                let provider_stats = stats.entry(provider.clone()).or_default();
121                let tool_entry = provider_stats
122                    .entry(tool_call.name().to_string())
123                    .or_insert((0, None));
124
125                tool_entry.0 += 1;
126
127                if tool_entry.1.is_none() {
128                    let result = tool_results.get(&event.id).cloned();
129                    // Serialize entire tool_call to extract arguments field
130                    let arguments = serde_json::to_value(tool_call)
131                        .ok()
132                        .and_then(|v| v.get("arguments").cloned())
133                        .and_then(|v| serde_json::to_string(&v).ok())
134                        .unwrap_or_else(|| String::from("(failed to serialize)"));
135                    tool_entry.1 = Some(ToolSample { arguments, result });
136                }
137            }
138        }
139    }
140
141    let provider_stats: ProviderStats = stats
142        .into_iter()
143        .map(|(provider_name, tools)| {
144            let sorted_tools: BTreeMap<_, _> = tools.into_iter().collect();
145
146            let classifications: Vec<ToolInfo> = sorted_tools
147                .keys()
148                .map(|tool_name| {
149                    let (origin, kind) = if let Ok(adapter) = create_adapter(&provider_name) {
150                        let (o, k) = adapter.mapper.classify(tool_name);
151                        (Some(format!("{:?}", o)), Some(format!("{:?}", k)))
152                    } else {
153                        (None, None)
154                    };
155
156                    ToolInfo {
157                        tool_name: tool_name.clone(),
158                        origin,
159                        kind,
160                    }
161                })
162                .collect();
163
164            (provider_name, (sorted_tools, classifications))
165        })
166        .collect();
167
168    Ok(StatsResult {
169        total_sessions,
170        provider_stats,
171    })
172}
agtrace_runtime/ops/analyze.rs

agtrace_runtime/ops/
analyze.rs
