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

agtrace_runtime/ops/
analyze.rs
