embacle/
metrics.rs

1// ABOUTME: Decorator wrapping any LlmProvider to measure latency, token usage, cost, and call counts
2// ABOUTME: Provides MetricsReport snapshots for cost and performance normalization, optional OTel export
3//
4// SPDX-License-Identifier: Apache-2.0
5// Copyright (c) 2026 dravr.ai
6
7//! # Cost/Latency Normalization
8//!
9//! [`MetricsProvider`] is a decorator that wraps any `Box<dyn LlmProvider>`,
10//! measuring per-call latency and token usage. Callers can retrieve a
11//! [`MetricsReport`] snapshot at any time via [`MetricsProvider::report()`].
12//!
13//! ## Cost Tracking
14//!
15//! Every call is priced through [`crate::pricing`], keyed on the wrapped
16//! provider's `name()` and the response's model, with the cache and
17//! reasoning counts the provider reported. A pair the table does not price
18//! costs `0.0` (see [`crate::pricing::calculate_cost_for`] for how that is
19//! logged). Costs accumulate in the report.
20//!
21//! ## OpenTelemetry (feature: `otel`)
22//!
23//! When built with `--features otel`, instruments are created via the
24//! global meter `embacle` and recorded on each `complete()` call.
25//!
26//! Token estimation: when `TokenUsage` is not provided by the inner provider,
27//! tokens are estimated at ~4 characters per token.
28
29use std::sync::{Arc, Mutex};
30use std::time::Instant;
31
32#[cfg(feature = "otel")]
33use opentelemetry::global;
34#[cfg(feature = "otel")]
35use opentelemetry::metrics::{Counter, Histogram};
36
37use async_trait::async_trait;
38use tracing::info;
39
40use crate::pricing::{calculate_cost_for, TokenCounts};
41use crate::types::{
42    ChatRequest, ChatResponse, ChatStream, LlmCapabilities, LlmProvider, RunnerError, TokenUsage,
43    ToolStream,
44};
45
46/// Characters-per-token estimate used when the provider does not report usage
47const CHARS_PER_TOKEN_ESTIMATE: u32 = 4;
48
49/// Accumulated metrics state protected by a mutex
50#[derive(Debug, Default)]
51struct MetricsState {
52    call_count: u64,
53    total_latency_ms: u64,
54    total_prompt_tokens: u64,
55    total_completion_tokens: u64,
56    total_tokens: u64,
57    errors_count: u64,
58    total_cost: f64,
59}
60
61/// Snapshot of accumulated metrics for a provider
62#[derive(Debug, Clone)]
63pub struct MetricsReport {
64    /// Name of the wrapped provider
65    pub provider_name: String,
66    /// Total number of `complete()` calls
67    pub call_count: u64,
68    /// Total latency across all calls (milliseconds)
69    pub total_latency_ms: u64,
70    /// Average latency per call (milliseconds)
71    pub avg_latency_ms: u64,
72    /// Total prompt tokens consumed
73    pub total_prompt_tokens: u64,
74    /// Total completion tokens generated
75    pub total_completion_tokens: u64,
76    /// Total tokens (prompt + completion)
77    pub total_tokens: u64,
78    /// Number of calls that returned an error
79    pub errors_count: u64,
80    /// Accumulated cost in USD (0.0 for every call the price table does not cover)
81    pub total_cost: f64,
82}
83
84/// OpenTelemetry instruments for metrics export
85#[cfg(feature = "otel")]
86struct OtelInstruments {
87    requests_total: Counter<u64>,
88    requests_duration_ms: Histogram<f64>,
89    tokens_prompt: Counter<u64>,
90    tokens_completion: Counter<u64>,
91    errors_total: Counter<u64>,
92    cost_total: Counter<f64>,
93}
94
95#[cfg(feature = "otel")]
96impl OtelInstruments {
97    fn new() -> Self {
98        let meter = global::meter("embacle");
99        Self {
100            requests_total: meter
101                .u64_counter("embacle.requests.total")
102                .with_description("Total LLM requests")
103                .build(),
104            requests_duration_ms: meter
105                .f64_histogram("embacle.requests.duration_ms")
106                .with_description("Request duration in milliseconds")
107                .build(),
108            tokens_prompt: meter
109                .u64_counter("embacle.tokens.prompt")
110                .with_description("Total prompt tokens consumed")
111                .build(),
112            tokens_completion: meter
113                .u64_counter("embacle.tokens.completion")
114                .with_description("Total completion tokens generated")
115                .build(),
116            errors_total: meter
117                .u64_counter("embacle.errors.total")
118                .with_description("Total error count")
119                .build(),
120            cost_total: meter
121                .f64_counter("embacle.cost.total")
122                .with_description("Total cost in USD")
123                .build(),
124        }
125    }
126}
127
128/// Decorator wrapping any `Box<dyn LlmProvider>` to collect latency, token, and cost metrics.
129///
130/// # Usage
131///
132/// ```rust,no_run
133/// # use embacle::metrics::MetricsProvider;
134/// # use embacle::types::LlmProvider;
135/// # fn example(provider: Box<dyn LlmProvider>) {
136/// let metered = MetricsProvider::new(provider);
137/// // ... use metered as LlmProvider ...
138/// let report = metered.report().unwrap(); // Safe: test assertion
139/// println!("calls={} avg_latency={}ms", report.call_count, report.avg_latency_ms);
140/// # }
141/// ```
142pub struct MetricsProvider {
143    inner: Box<dyn LlmProvider>,
144    state: Arc<Mutex<MetricsState>>,
145    #[cfg(feature = "otel")]
146    otel: OtelInstruments,
147}
148
149impl MetricsProvider {
150    /// Wrap a provider with metrics collection
151    pub fn new(inner: Box<dyn LlmProvider>) -> Self {
152        Self {
153            inner,
154            state: Arc::new(Mutex::new(MetricsState::default())),
155            #[cfg(feature = "otel")]
156            otel: OtelInstruments::new(),
157        }
158    }
159
160    /// Return a snapshot of the current metrics
161    ///
162    /// # Errors
163    ///
164    /// Returns [`RunnerError`] if the internal mutex is poisoned.
165    pub fn report(&self) -> Result<MetricsReport, RunnerError> {
166        let state = self
167            .state
168            .lock()
169            .map_err(|_| RunnerError::internal("metrics lock poisoned"))?;
170        let divisor = state.call_count.max(1);
171        Ok(MetricsReport {
172            provider_name: self.inner.name().to_owned(),
173            call_count: state.call_count,
174            total_latency_ms: state.total_latency_ms,
175            avg_latency_ms: state.total_latency_ms / divisor,
176            total_prompt_tokens: state.total_prompt_tokens,
177            total_completion_tokens: state.total_completion_tokens,
178            total_tokens: state.total_tokens,
179            errors_count: state.errors_count,
180            total_cost: state.total_cost,
181        })
182    }
183
184    /// Reset all counters to zero
185    ///
186    /// # Errors
187    ///
188    /// Returns [`RunnerError`] if the internal mutex is poisoned.
189    pub fn reset(&self) -> Result<(), RunnerError> {
190        let mut state = self
191            .state
192            .lock()
193            .map_err(|_| RunnerError::internal("metrics lock poisoned"))?;
194        *state = MetricsState::default();
195        Ok(())
196    }
197
198    /// Price one call: the wrapped provider's `name()` is the table key, and
199    /// the cache and reasoning counts the provider reported ride along.
200    /// `prompt_tokens`/`completion_tokens` are the (possibly estimated)
201    /// headline counts; `usage` supplies the optional categories.
202    fn compute_cost(
203        &self,
204        model: &str,
205        prompt_tokens: u64,
206        completion_tokens: u64,
207        usage: Option<&TokenUsage>,
208    ) -> f64 {
209        let count = |n: u64| i64::try_from(n).unwrap_or(i64::MAX);
210        let optional = |n: Option<u32>| i64::from(n.unwrap_or(0));
211        let counts = TokenCounts::new(count(prompt_tokens), count(completion_tokens))
212            .with_cache(
213                optional(usage.and_then(|u| u.cached_read_tokens)),
214                optional(usage.and_then(|u| u.cached_write_tokens)),
215            )
216            .with_reasoning(optional(usage.and_then(|u| u.reasoning_tokens)));
217        calculate_cost_for(self.inner.name(), model, &counts)
218    }
219}
220
221/// Estimate token count from character length (~4 chars per token)
222fn estimate_tokens(text: &str) -> u32 {
223    #[allow(clippy::cast_possible_truncation)]
224    let len = text.len() as u32;
225    len / CHARS_PER_TOKEN_ESTIMATE.max(1)
226}
227
228#[async_trait]
229impl LlmProvider for MetricsProvider {
230    fn name(&self) -> &'static str {
231        self.inner.name()
232    }
233
234    fn display_name(&self) -> &str {
235        self.inner.display_name()
236    }
237
238    fn capabilities(&self) -> LlmCapabilities {
239        self.inner.capabilities()
240    }
241
242    fn default_model(&self) -> &str {
243        self.inner.default_model()
244    }
245
246    fn available_models(&self) -> &[String] {
247        self.inner.available_models()
248    }
249
250    async fn complete(&self, request: &ChatRequest) -> Result<ChatResponse, RunnerError> {
251        let start = Instant::now();
252        let result = self.inner.complete(request).await;
253        #[allow(clippy::cast_possible_truncation)]
254        let elapsed_ms = start.elapsed().as_millis() as u64;
255
256        let mut state = self
257            .state
258            .lock()
259            .map_err(|_| RunnerError::internal("metrics lock poisoned"))?;
260        state.call_count += 1;
261        state.total_latency_ms += elapsed_ms;
262
263        #[cfg(feature = "otel")]
264        let otel_attrs: Vec<opentelemetry::KeyValue> = {
265            let mut attrs = vec![opentelemetry::KeyValue::new("provider", self.inner.name())];
266            if let Some(turn_id) = request.turn_id {
267                attrs.push(opentelemetry::KeyValue::new("turn_id", turn_id.to_string()));
268            }
269            attrs
270        };
271
272        if let Ok(response) = &result {
273            let usage = response.usage.as_ref();
274            let prompt_tokens = u64::from(
275                usage.map_or_else(|| estimate_prompt_tokens(request), |u| u.prompt_tokens),
276            );
277            let completion_tokens = u64::from(usage.map_or_else(
278                || estimate_tokens(&response.content),
279                |u| u.completion_tokens,
280            ));
281            let total = prompt_tokens + completion_tokens;
282
283            state.total_prompt_tokens += prompt_tokens;
284            state.total_completion_tokens += completion_tokens;
285            state.total_tokens += total;
286
287            let cost = self.compute_cost(&response.model, prompt_tokens, completion_tokens, usage);
288            state.total_cost += cost;
289
290            info!(
291                provider = self.inner.name(),
292                elapsed_ms, prompt_tokens, completion_tokens, cost, "metrics: complete() succeeded"
293            );
294
295            #[cfg(feature = "otel")]
296            {
297                let attrs = otel_attrs.as_slice();
298                self.otel.requests_total.add(1, attrs);
299                #[allow(clippy::cast_precision_loss)]
300                self.otel
301                    .requests_duration_ms
302                    .record(elapsed_ms as f64, attrs);
303                self.otel.tokens_prompt.add(prompt_tokens, attrs);
304                self.otel.tokens_completion.add(completion_tokens, attrs);
305                if cost > 0.0 {
306                    self.otel.cost_total.add(cost, attrs);
307                }
308            }
309        } else {
310            state.errors_count += 1;
311            info!(
312                provider = self.inner.name(),
313                elapsed_ms, "metrics: complete() failed"
314            );
315
316            #[cfg(feature = "otel")]
317            {
318                self.otel.errors_total.add(1, otel_attrs.as_slice());
319            }
320        }
321
322        drop(state);
323
324        result
325    }
326
327    /// Delegate streaming directly; only measures stream setup time (documented limitation)
328    async fn complete_stream(&self, request: &ChatRequest) -> Result<ChatStream, RunnerError> {
329        self.inner.complete_stream(request).await
330    }
331
332    async fn complete_stream_with_tools(
333        &self,
334        request: &ChatRequest,
335    ) -> Result<ToolStream, RunnerError> {
336        self.inner.complete_stream_with_tools(request).await
337    }
338
339    async fn health_check(&self) -> Result<bool, RunnerError> {
340        self.inner.health_check().await
341    }
342}
343
344/// Estimate prompt tokens from request messages
345fn estimate_prompt_tokens(request: &ChatRequest) -> u32 {
346    let total_chars: usize = request.messages.iter().map(|m| m.content.len()).sum();
347    #[allow(clippy::cast_possible_truncation)]
348    let len = total_chars as u32;
349    len / CHARS_PER_TOKEN_ESTIMATE.max(1)
350}
351
352#[cfg(test)]
353mod tests {
354    use super::*;
355    use crate::types::{
356        ChatMessage, ChatRequest, ChatResponse, ChatStream, LlmCapabilities, LlmProvider,
357        RunnerError, TokenUsage,
358    };
359    use async_trait::async_trait;
360    use std::sync::atomic::{AtomicU32, Ordering};
361
362    struct TestProvider {
363        provider_name: &'static str,
364        responses: Mutex<Vec<Result<ChatResponse, RunnerError>>>,
365        call_count: AtomicU32,
366    }
367
368    impl TestProvider {
369        fn new(responses: Vec<Result<ChatResponse, RunnerError>>) -> Self {
370            Self::named("test", responses)
371        }
372
373        /// A provider reporting `name`, so the price table keys on it.
374        fn named(name: &'static str, responses: Vec<Result<ChatResponse, RunnerError>>) -> Self {
375            Self {
376                provider_name: name,
377                responses: Mutex::new(responses),
378                call_count: AtomicU32::new(0),
379            }
380        }
381    }
382
383    #[async_trait]
384    impl LlmProvider for TestProvider {
385        fn name(&self) -> &'static str {
386            self.provider_name
387        }
388        fn display_name(&self) -> &str {
389            "Test Provider"
390        }
391        fn capabilities(&self) -> LlmCapabilities {
392            LlmCapabilities::text_only()
393        }
394        fn default_model(&self) -> &'static str {
395            "test-model"
396        }
397        fn available_models(&self) -> &[String] {
398            &[]
399        }
400
401        async fn complete(&self, _request: &ChatRequest) -> Result<ChatResponse, RunnerError> {
402            self.call_count.fetch_add(1, Ordering::SeqCst);
403            let mut responses = self.responses.lock().expect("test lock"); // Safe: test assertion
404            if responses.is_empty() {
405                Ok(ChatResponse {
406                    content: "default".to_owned(),
407                    model: "test-model".to_owned(),
408                    usage: None,
409                    finish_reason: Some("stop".to_owned()),
410                    warnings: None,
411                    tool_calls: None,
412                })
413            } else {
414                responses.remove(0)
415            }
416        }
417
418        async fn complete_stream(&self, _request: &ChatRequest) -> Result<ChatStream, RunnerError> {
419            Err(RunnerError::internal("streaming not supported in test"))
420        }
421
422        async fn health_check(&self) -> Result<bool, RunnerError> {
423            Ok(true)
424        }
425    }
426
427    #[test]
428    fn fresh_report_is_zeroed() {
429        let provider = TestProvider::new(vec![]);
430        let metered = MetricsProvider::new(Box::new(provider));
431        let report = metered.report().unwrap(); // Safe: test assertion
432        assert_eq!(report.call_count, 0);
433        assert_eq!(report.total_latency_ms, 0);
434        assert_eq!(report.avg_latency_ms, 0);
435        assert_eq!(report.total_prompt_tokens, 0);
436        assert_eq!(report.total_completion_tokens, 0);
437        assert_eq!(report.total_tokens, 0);
438        assert_eq!(report.errors_count, 0);
439        assert!(
440            report.total_cost.abs() < f64::EPSILON,
441            "{}",
442            report.total_cost
443        );
444        assert_eq!(report.provider_name, "test");
445    }
446
447    #[tokio::test]
448    async fn call_count_increments() {
449        let provider = TestProvider::new(vec![
450            Ok(ChatResponse {
451                content: "hello world".to_owned(),
452                model: "test-model".to_owned(),
453                usage: Some(TokenUsage::new(10, 5, 15)),
454                finish_reason: Some("stop".to_owned()),
455                warnings: None,
456                tool_calls: None,
457            }),
458            Ok(ChatResponse {
459                content: "second".to_owned(),
460                model: "test-model".to_owned(),
461                usage: Some(TokenUsage::new(8, 3, 11)),
462                finish_reason: Some("stop".to_owned()),
463                warnings: None,
464                tool_calls: None,
465            }),
466        ]);
467        let metered = MetricsProvider::new(Box::new(provider));
468        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
469
470        metered.complete(&request).await.expect("first call"); // Safe: test assertion
471        metered.complete(&request).await.expect("second call"); // Safe: test assertion
472
473        let report = metered.report().unwrap(); // Safe: test assertion
474        assert_eq!(report.call_count, 2);
475        assert_eq!(report.total_prompt_tokens, 18);
476        assert_eq!(report.total_completion_tokens, 8);
477        assert_eq!(report.total_tokens, 26);
478        assert_eq!(report.errors_count, 0);
479    }
480
481    #[tokio::test]
482    async fn errors_count_on_failure() {
483        let provider = TestProvider::new(vec![Err(RunnerError::external_service("test", "boom"))]);
484        let metered = MetricsProvider::new(Box::new(provider));
485        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
486
487        let result = metered.complete(&request).await;
488        assert!(result.is_err());
489
490        let report = metered.report().unwrap(); // Safe: test assertion
491        assert_eq!(report.call_count, 1);
492        assert_eq!(report.errors_count, 1);
493    }
494
495    #[tokio::test]
496    async fn token_estimation_when_no_usage() {
497        let provider = TestProvider::new(vec![Ok(ChatResponse {
498            content: "abcdefghijklmnop".to_owned(), // 16 chars => 4 tokens
499            model: "test-model".to_owned(),
500            usage: None,
501            finish_reason: Some("stop".to_owned()),
502            warnings: None,
503            tool_calls: None,
504        })]);
505        let metered = MetricsProvider::new(Box::new(provider));
506        let request = ChatRequest::new(vec![ChatMessage::user("12345678")]); // 8 chars => 2 tokens
507
508        metered.complete(&request).await.expect("call"); // Safe: test assertion
509
510        let report = metered.report().unwrap(); // Safe: test assertion
511        assert_eq!(report.total_prompt_tokens, 2);
512        assert_eq!(report.total_completion_tokens, 4);
513        assert_eq!(report.total_tokens, 6);
514    }
515
516    #[test]
517    fn div_by_zero_guard_on_avg_latency() {
518        let provider = TestProvider::new(vec![]);
519        let metered = MetricsProvider::new(Box::new(provider));
520        // No calls made — avg should be 0/max(0,1) = 0, not panic
521        let report = metered.report().unwrap(); // Safe: test assertion
522        assert_eq!(report.avg_latency_ms, 0);
523    }
524
525    #[tokio::test]
526    async fn reset_zeroes_counters() {
527        let provider = TestProvider::new(vec![Ok(ChatResponse {
528            content: "hello".to_owned(),
529            model: "test-model".to_owned(),
530            usage: Some(TokenUsage::new(5, 2, 7)),
531            finish_reason: Some("stop".to_owned()),
532            warnings: None,
533            tool_calls: None,
534        })]);
535        let metered = MetricsProvider::new(Box::new(provider));
536        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
537
538        metered.complete(&request).await.expect("call"); // Safe: test assertion
539        assert_eq!(metered.report().unwrap().call_count, 1); // Safe: test assertion
540
541        metered.reset().unwrap(); // Safe: test assertion
542        let report = metered.report().unwrap(); // Safe: test assertion
543        assert_eq!(report.call_count, 0);
544        assert_eq!(report.total_tokens, 0);
545        assert_eq!(report.errors_count, 0);
546        assert!(
547            report.total_cost.abs() < f64::EPSILON,
548            "{}",
549            report.total_cost
550        );
551    }
552
553    // ========================================================================
554    // Cost tracking tests — priced through `crate::pricing`, keyed on name()
555    // ========================================================================
556
557    /// gemini-2.5-flash: $0.15/M input, $0.60/M output.
558    const FLASH_IN: f64 = 0.15 / 1_000_000.0;
559    const FLASH_OUT: f64 = 0.60 / 1_000_000.0;
560
561    fn flash_response(content: &str, usage: Option<TokenUsage>) -> ChatResponse {
562        ChatResponse {
563            content: content.to_owned(),
564            model: "gemini-2.5-flash".to_owned(),
565            usage,
566            finish_reason: Some("stop".to_owned()),
567            warnings: None,
568            tool_calls: None,
569        }
570    }
571
572    #[tokio::test]
573    async fn cost_with_known_model() {
574        let provider = TestProvider::named(
575            "gemini",
576            vec![Ok(flash_response(
577                "response",
578                Some(TokenUsage::new(1000, 500, 1500)),
579            ))],
580        );
581        let metered = MetricsProvider::new(Box::new(provider));
582        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
583        metered.complete(&request).await.expect("call"); // Safe: test assertion
584
585        let report = metered.report().unwrap(); // Safe: test assertion
586        let expected = 1000.0f64.mul_add(FLASH_IN, 500.0 * FLASH_OUT);
587        assert!(
588            (report.total_cost - expected).abs() < 1e-12,
589            "{}",
590            report.total_cost
591        );
592    }
593
594    /// The table is keyed on the provider's `name()`: the same model under a
595    /// provider the table does not price costs nothing.
596    #[tokio::test]
597    async fn cost_with_unknown_provider_is_zero() {
598        let provider = TestProvider::new(vec![Ok(flash_response(
599            "response",
600            Some(TokenUsage::new(1000, 500, 1500)),
601        ))]);
602        let metered = MetricsProvider::new(Box::new(provider));
603        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
604        metered.complete(&request).await.expect("call"); // Safe: test assertion
605
606        let report = metered.report().unwrap(); // Safe: test assertion
607        assert!(
608            report.total_cost.abs() < f64::EPSILON,
609            "{}",
610            report.total_cost
611        );
612    }
613
614    #[tokio::test]
615    async fn cost_with_unknown_model() {
616        let provider = TestProvider::named(
617            "gemini",
618            vec![Ok(ChatResponse {
619                content: "response".to_owned(),
620                model: "some-unknown-model".to_owned(),
621                usage: Some(TokenUsage::new(1000, 500, 1500)),
622                finish_reason: Some("stop".to_owned()),
623                warnings: None,
624                tool_calls: None,
625            })],
626        );
627        let metered = MetricsProvider::new(Box::new(provider));
628        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
629        metered.complete(&request).await.expect("call"); // Safe: test assertion
630
631        let report = metered.report().unwrap(); // Safe: test assertion
632        assert!(
633            report.total_cost.abs() < f64::EPSILON,
634            "{}",
635            report.total_cost
636        );
637    }
638
639    #[tokio::test]
640    async fn cost_accumulates() {
641        let provider = TestProvider::named(
642            "gemini",
643            vec![
644                Ok(flash_response("r1", Some(TokenUsage::new(1000, 500, 1500)))),
645                Ok(flash_response(
646                    "r2",
647                    Some(TokenUsage::new(2000, 1000, 3000)),
648                )),
649            ],
650        );
651        let metered = MetricsProvider::new(Box::new(provider));
652        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
653        metered.complete(&request).await.expect("call 1"); // Safe: test assertion
654        metered.complete(&request).await.expect("call 2"); // Safe: test assertion
655
656        let report = metered.report().unwrap(); // Safe: test assertion
657        let expected = 3000.0f64.mul_add(FLASH_IN, 1500.0 * FLASH_OUT);
658        assert!(
659            (report.total_cost - expected).abs() < 1e-12,
660            "{}",
661            report.total_cost
662        );
663    }
664
665    /// Cache reads the provider reports bill at the model's read multiplier
666    /// (Gemini 0.25×), not as fresh input.
667    #[tokio::test]
668    async fn cost_credits_reported_cache_reads() {
669        let provider = TestProvider::named(
670            "gemini",
671            vec![Ok(flash_response(
672                "r",
673                Some(TokenUsage::new(1000, 0, 1000).with_cache(Some(1000), None)),
674            ))],
675        );
676        let metered = MetricsProvider::new(Box::new(provider));
677        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
678        metered.complete(&request).await.expect("call"); // Safe: test assertion
679
680        let report = metered.report().unwrap(); // Safe: test assertion
681        let expected = 1000.0 * FLASH_IN * 0.25;
682        assert!(
683            (report.total_cost - expected).abs() < 1e-12,
684            "{}",
685            report.total_cost
686        );
687    }
688
689    #[tokio::test]
690    async fn cost_with_estimated_tokens() {
691        // 16 chars => 4 completion tokens; the request's 8 chars => 2 prompt tokens
692        let provider =
693            TestProvider::named("gemini", vec![Ok(flash_response("abcdefghijklmnop", None))]);
694        let metered = MetricsProvider::new(Box::new(provider));
695        let request = ChatRequest::new(vec![ChatMessage::user("12345678")]);
696        metered.complete(&request).await.expect("call"); // Safe: test assertion
697
698        let report = metered.report().unwrap(); // Safe: test assertion
699        let expected = 2.0f64.mul_add(FLASH_IN, 4.0 * FLASH_OUT);
700        assert!(report.total_cost > 0.0);
701        assert!(
702            (report.total_cost - expected).abs() < 1e-12,
703            "{}",
704            report.total_cost
705        );
706    }
707
708    #[tokio::test]
709    async fn reset_zeroes_cost() {
710        let provider = TestProvider::named(
711            "gemini",
712            vec![Ok(flash_response(
713                "response",
714                Some(TokenUsage::new(1000, 500, 1500)),
715            ))],
716        );
717        let metered = MetricsProvider::new(Box::new(provider));
718        let request = ChatRequest::new(vec![ChatMessage::user("hi")]);
719        metered.complete(&request).await.expect("call"); // Safe: test assertion
720        assert!(metered.report().unwrap().total_cost > 0.0); // Safe: test assertion
721
722        metered.reset().unwrap(); // Safe: test assertion
723        let cost = metered.report().unwrap().total_cost; // Safe: test assertion
724        assert!(cost.abs() < f64::EPSILON, "{cost}");
725    }
726}
embacle/metrics.rs

embacle/
metrics.rs
