kova-sdk 0.1.0

Async-first Rust library for building LLM-powered agents with tool calling, streaming, MCP, and multi-agent orchestration
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217

//! Metrics collection for LLM requests, tool invocations, and token usage.
//!
//! `MetricsCollector` uses atomic counters and `RwLock`-protected histograms
//! so it works both with and without the `telemetry` feature flag.
//! It is `Send + Sync` by construction.

use std::sync::RwLock;
use std::sync::atomic::{AtomicU64, Ordering};

/// Collects quantitative measurements for observability.
///
/// Records histograms for LLM request latency and tool execution duration,
/// and counters for total LLM requests, tool invocations, tokens consumed,
/// and errors.
#[derive(Debug)]
pub struct MetricsCollector {
    total_llm_requests: AtomicU64,
    total_tool_invocations: AtomicU64,
    total_tokens_consumed: AtomicU64,
    total_errors: AtomicU64,
    llm_request_latency_ms: RwLock<Vec<f64>>,
    tool_execution_duration_ms: RwLock<Vec<f64>>,
}

impl MetricsCollector {
    /// Create a new, zeroed `MetricsCollector`.
    pub fn new() -> Self {
        Self {
            total_llm_requests: AtomicU64::new(0),
            total_tool_invocations: AtomicU64::new(0),
            total_tokens_consumed: AtomicU64::new(0),
            total_errors: AtomicU64::new(0),
            llm_request_latency_ms: RwLock::new(Vec::new()),
            tool_execution_duration_ms: RwLock::new(Vec::new()),
        }
    }

    /// Record a successful LLM request.
    ///
    /// Increments the LLM request counter, adds `latency_ms` to the latency
    /// histogram, and accumulates `input_tokens + output_tokens` into the
    /// total token counter.
    pub fn record_llm_request(&self, latency_ms: f64, input_tokens: u64, output_tokens: u64) {
        self.total_llm_requests.fetch_add(1, Ordering::Relaxed);
        self.total_tokens_consumed
            .fetch_add(input_tokens + output_tokens, Ordering::Relaxed);
        if let Ok(mut hist) = self.llm_request_latency_ms.write() {
            hist.push(latency_ms);
        }
    }

    /// Record an LLM error (increments the error counter).
    pub fn record_llm_error(&self) {
        self.total_errors.fetch_add(1, Ordering::Relaxed);
    }

    /// Record a tool invocation.
    ///
    /// Increments the tool invocation counter, adds `duration_ms` to the
    /// duration histogram, and increments the error counter when `success`
    /// is `false`.
    pub fn record_tool_invocation(&self, duration_ms: f64, success: bool) {
        self.total_tool_invocations.fetch_add(1, Ordering::Relaxed);
        if !success {
            self.total_errors.fetch_add(1, Ordering::Relaxed);
        }
        if let Ok(mut hist) = self.tool_execution_duration_ms.write() {
            hist.push(duration_ms);
        }
    }

    /// Total number of LLM requests recorded.
    pub fn llm_request_count(&self) -> u64 {
        self.total_llm_requests.load(Ordering::Relaxed)
    }

    /// Total number of tool invocations recorded.
    pub fn tool_invocation_count(&self) -> u64 {
        self.total_tool_invocations.load(Ordering::Relaxed)
    }

    /// Total tokens consumed (input + output) across all LLM requests.
    pub fn total_tokens(&self) -> u64 {
        self.total_tokens_consumed.load(Ordering::Relaxed)
    }

    /// Total error count (LLM errors + failed tool invocations).
    pub fn error_count(&self) -> u64 {
        self.total_errors.load(Ordering::Relaxed)
    }

    /// Clone of the recorded LLM request latencies (milliseconds).
    pub fn llm_latency_histogram(&self) -> Vec<f64> {
        self.llm_request_latency_ms
            .read()
            .map(|h| h.clone())
            .unwrap_or_default()
    }

    /// Clone of the recorded tool execution durations (milliseconds).
    pub fn tool_duration_histogram(&self) -> Vec<f64> {
        self.tool_execution_duration_ms
            .read()
            .map(|h| h.clone())
            .unwrap_or_default()
    }

    /// Reset all counters and histograms to zero / empty.
    pub fn reset(&self) {
        self.total_llm_requests.store(0, Ordering::Relaxed);
        self.total_tool_invocations.store(0, Ordering::Relaxed);
        self.total_tokens_consumed.store(0, Ordering::Relaxed);
        self.total_errors.store(0, Ordering::Relaxed);
        if let Ok(mut hist) = self.llm_request_latency_ms.write() {
            hist.clear();
        }
        if let Ok(mut hist) = self.tool_execution_duration_ms.write() {
            hist.clear();
        }
    }
}

impl Default for MetricsCollector {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn new_collector_starts_at_zero() {
        let mc = MetricsCollector::new();
        assert_eq!(mc.llm_request_count(), 0);
        assert_eq!(mc.tool_invocation_count(), 0);
        assert_eq!(mc.total_tokens(), 0);
        assert_eq!(mc.error_count(), 0);
        assert!(mc.llm_latency_histogram().is_empty());
        assert!(mc.tool_duration_histogram().is_empty());
    }

    #[test]
    fn record_llm_request_increments_counters_and_histogram() {
        let mc = MetricsCollector::new();
        mc.record_llm_request(42.5, 100, 50);
        mc.record_llm_request(10.0, 200, 80);

        assert_eq!(mc.llm_request_count(), 2);
        assert_eq!(mc.total_tokens(), 430); // 100+50 + 200+80
        assert_eq!(mc.llm_latency_histogram(), vec![42.5, 10.0]);
    }

    #[test]
    fn record_tool_invocation_increments_counters() {
        let mc = MetricsCollector::new();
        mc.record_tool_invocation(5.0, true);
        mc.record_tool_invocation(12.3, true);
        mc.record_tool_invocation(8.0, false);

        assert_eq!(mc.tool_invocation_count(), 3);
        assert_eq!(mc.tool_duration_histogram(), vec![5.0, 12.3, 8.0]);
    }

    #[test]
    fn error_counting_for_llm_and_tool() {
        let mc = MetricsCollector::new();
        mc.record_llm_error();
        mc.record_llm_error();
        mc.record_tool_invocation(1.0, false);

        assert_eq!(mc.error_count(), 3);
    }

    #[test]
    fn token_accumulation() {
        let mc = MetricsCollector::new();
        mc.record_llm_request(1.0, 10, 5);
        mc.record_llm_request(1.0, 20, 10);
        mc.record_llm_request(1.0, 0, 0);

        assert_eq!(mc.total_tokens(), 45); // 15 + 30 + 0
    }

    #[test]
    fn reset_clears_everything() {
        let mc = MetricsCollector::new();
        mc.record_llm_request(50.0, 100, 200);
        mc.record_tool_invocation(10.0, true);
        mc.record_llm_error();
        mc.record_tool_invocation(5.0, false);

        // Sanity: non-zero before reset
        assert!(mc.llm_request_count() > 0);
        assert!(mc.tool_invocation_count() > 0);
        assert!(mc.total_tokens() > 0);
        assert!(mc.error_count() > 0);
        assert!(!mc.llm_latency_histogram().is_empty());
        assert!(!mc.tool_duration_histogram().is_empty());

        mc.reset();

        assert_eq!(mc.llm_request_count(), 0);
        assert_eq!(mc.tool_invocation_count(), 0);
        assert_eq!(mc.total_tokens(), 0);
        assert_eq!(mc.error_count(), 0);
        assert!(mc.llm_latency_histogram().is_empty());
        assert!(mc.tool_duration_histogram().is_empty());
    }

    #[test]
    fn metrics_collector_is_send_and_sync() {
        fn assert_send_sync<T: Send + Sync>() {}
        assert_send_sync::<MetricsCollector>();
    }
}