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

agtrace_runtime/ops/
analyze.rs
