litellm_rs/sdk/client/
types.rs

1//! Type definitions for the LLM client
2
3use dashmap::DashMap;
4use std::sync::atomic::AtomicUsize;
5use std::time::SystemTime;
6
7/// Provider statistics
8#[derive(Debug, Clone, Default)]
9pub struct ProviderStats {
10    pub requests: u64,
11    pub errors: u64,
12    pub total_tokens: u64,
13    pub total_cost: f64,
14    pub avg_latency_ms: f64,
15    pub last_used: Option<SystemTime>,
16    pub health_score: f64,
17}
18
19/// Load balancer
20#[derive(Debug)]
21pub struct LoadBalancer {
22    pub(crate) strategy: LoadBalancingStrategy,
23    pub(crate) round_robin_counters: DashMap<String, AtomicUsize>,
24}
25
26/// Load balancing strategy
27#[derive(Debug, Clone)]
28pub enum LoadBalancingStrategy {
29    RoundRobin,
30    LeastLatency,
31    WeightedRandom,
32    HealthBased,
33}
34
35impl LoadBalancer {
36    pub(crate) fn new(strategy: LoadBalancingStrategy) -> Self {
37        Self {
38            strategy,
39            round_robin_counters: DashMap::new(),
40        }
41    }
42}
43
44#[cfg(test)]
45mod tests {
46    use super::*;
47
48    // ==================== ProviderStats Tests ====================
49
50    #[test]
51    fn test_provider_stats_default() {
52        let stats = ProviderStats::default();
53
54        assert_eq!(stats.requests, 0);
55        assert_eq!(stats.errors, 0);
56        assert_eq!(stats.total_tokens, 0);
57        assert!((stats.total_cost - 0.0).abs() < f64::EPSILON);
58        assert!((stats.avg_latency_ms - 0.0).abs() < f64::EPSILON);
59        assert!(stats.last_used.is_none());
60        assert!((stats.health_score - 0.0).abs() < f64::EPSILON);
61    }
62
63    #[test]
64    fn test_provider_stats_with_values() {
65        let stats = ProviderStats {
66            requests: 1000,
67            errors: 10,
68            total_tokens: 500000,
69            total_cost: 25.50,
70            avg_latency_ms: 150.0,
71            last_used: Some(SystemTime::now()),
72            health_score: 0.95,
73        };
74
75        assert_eq!(stats.requests, 1000);
76        assert_eq!(stats.errors, 10);
77        assert_eq!(stats.total_tokens, 500000);
78        assert!((stats.total_cost - 25.50).abs() < f64::EPSILON);
79    }
80
81    #[test]
82    fn test_provider_stats_clone() {
83        let stats = ProviderStats {
84            requests: 500,
85            errors: 5,
86            total_tokens: 250000,
87            total_cost: 12.75,
88            avg_latency_ms: 100.0,
89            last_used: None,
90            health_score: 0.90,
91        };
92
93        let cloned = stats.clone();
94        assert_eq!(cloned.requests, stats.requests);
95        assert_eq!(cloned.errors, stats.errors);
96        assert_eq!(cloned.health_score, stats.health_score);
97    }
98
99    #[test]
100    fn test_provider_stats_debug() {
101        let stats = ProviderStats::default();
102        let debug_str = format!("{:?}", stats);
103
104        assert!(debug_str.contains("ProviderStats"));
105        assert!(debug_str.contains("requests"));
106    }
107
108    #[test]
109    fn test_provider_stats_error_rate() {
110        let stats = ProviderStats {
111            requests: 100,
112            errors: 5,
113            total_tokens: 0,
114            total_cost: 0.0,
115            avg_latency_ms: 0.0,
116            last_used: None,
117            health_score: 0.0,
118        };
119
120        let error_rate = if stats.requests > 0 {
121            stats.errors as f64 / stats.requests as f64
122        } else {
123            0.0
124        };
125
126        assert!((error_rate - 0.05).abs() < f64::EPSILON);
127    }
128
129    #[test]
130    fn test_provider_stats_tokens_per_request() {
131        let stats = ProviderStats {
132            requests: 100,
133            errors: 0,
134            total_tokens: 50000,
135            total_cost: 0.0,
136            avg_latency_ms: 0.0,
137            last_used: None,
138            health_score: 0.0,
139        };
140
141        let avg_tokens = if stats.requests > 0 {
142            stats.total_tokens as f64 / stats.requests as f64
143        } else {
144            0.0
145        };
146
147        assert!((avg_tokens - 500.0).abs() < f64::EPSILON);
148    }
149
150    #[test]
151    fn test_provider_stats_cost_per_request() {
152        let stats = ProviderStats {
153            requests: 1000,
154            errors: 0,
155            total_tokens: 0,
156            total_cost: 10.0,
157            avg_latency_ms: 0.0,
158            last_used: None,
159            health_score: 0.0,
160        };
161
162        let cost_per_request = if stats.requests > 0 {
163            stats.total_cost / stats.requests as f64
164        } else {
165            0.0
166        };
167
168        assert!((cost_per_request - 0.01).abs() < f64::EPSILON);
169    }
170
171    #[test]
172    fn test_provider_stats_high_load() {
173        let stats = ProviderStats {
174            requests: 1_000_000,
175            errors: 1000,
176            total_tokens: 500_000_000,
177            total_cost: 5000.0,
178            avg_latency_ms: 200.0,
179            last_used: Some(SystemTime::now()),
180            health_score: 0.85,
181        };
182
183        assert!(stats.requests > 0);
184        assert!(stats.health_score > 0.8);
185    }
186
187    #[test]
188    fn test_provider_stats_unhealthy() {
189        let stats = ProviderStats {
190            requests: 100,
191            errors: 50, // 50% error rate
192            total_tokens: 10000,
193            total_cost: 1.0,
194            avg_latency_ms: 5000.0, // Very slow
195            last_used: Some(SystemTime::now()),
196            health_score: 0.2,
197        };
198
199        let error_rate = stats.errors as f64 / stats.requests as f64;
200        assert!(error_rate > 0.3);
201        assert!(stats.health_score < 0.5);
202    }
203
204    // ==================== LoadBalancingStrategy Tests ====================
205
206    #[test]
207    fn test_load_balancing_strategy_round_robin() {
208        let strategy = LoadBalancingStrategy::RoundRobin;
209        let debug_str = format!("{:?}", strategy);
210        assert!(debug_str.contains("RoundRobin"));
211    }
212
213    #[test]
214    fn test_load_balancing_strategy_least_latency() {
215        let strategy = LoadBalancingStrategy::LeastLatency;
216        let debug_str = format!("{:?}", strategy);
217        assert!(debug_str.contains("LeastLatency"));
218    }
219
220    #[test]
221    fn test_load_balancing_strategy_weighted_random() {
222        let strategy = LoadBalancingStrategy::WeightedRandom;
223        let debug_str = format!("{:?}", strategy);
224        assert!(debug_str.contains("WeightedRandom"));
225    }
226
227    #[test]
228    fn test_load_balancing_strategy_health_based() {
229        let strategy = LoadBalancingStrategy::HealthBased;
230        let debug_str = format!("{:?}", strategy);
231        assert!(debug_str.contains("HealthBased"));
232    }
233
234    #[test]
235    fn test_load_balancing_strategy_clone() {
236        let strategy = LoadBalancingStrategy::LeastLatency;
237        let cloned = strategy.clone();
238
239        let original_str = format!("{:?}", strategy);
240        let cloned_str = format!("{:?}", cloned);
241        assert_eq!(original_str, cloned_str);
242    }
243
244    // ==================== LoadBalancer Tests ====================
245
246    #[test]
247    fn test_load_balancer_new_round_robin() {
248        let balancer = LoadBalancer::new(LoadBalancingStrategy::RoundRobin);
249        let debug_str = format!("{:?}", balancer);
250        assert!(debug_str.contains("RoundRobin"));
251    }
252
253    #[test]
254    fn test_load_balancer_new_least_latency() {
255        let balancer = LoadBalancer::new(LoadBalancingStrategy::LeastLatency);
256        let debug_str = format!("{:?}", balancer);
257        assert!(debug_str.contains("LeastLatency"));
258    }
259
260    #[test]
261    fn test_load_balancer_new_weighted_random() {
262        let balancer = LoadBalancer::new(LoadBalancingStrategy::WeightedRandom);
263        let debug_str = format!("{:?}", balancer);
264        assert!(debug_str.contains("WeightedRandom"));
265    }
266
267    #[test]
268    fn test_load_balancer_new_health_based() {
269        let balancer = LoadBalancer::new(LoadBalancingStrategy::HealthBased);
270        let debug_str = format!("{:?}", balancer);
271        assert!(debug_str.contains("HealthBased"));
272    }
273
274    #[test]
275    fn test_load_balancer_debug() {
276        let balancer = LoadBalancer::new(LoadBalancingStrategy::RoundRobin);
277        let debug_str = format!("{:?}", balancer);
278        assert!(debug_str.contains("LoadBalancer"));
279    }
280
281    // ==================== Integration Tests ====================
282
283    #[test]
284    fn test_provider_selection_simulation() {
285        let providers = [
286            ProviderStats {
287                requests: 100,
288                errors: 1,
289                health_score: 0.99,
290                avg_latency_ms: 50.0,
291                ..ProviderStats::default()
292            },
293            ProviderStats {
294                requests: 200,
295                errors: 20,
296                health_score: 0.80,
297                avg_latency_ms: 100.0,
298                ..ProviderStats::default()
299            },
300            ProviderStats {
301                requests: 50,
302                errors: 0,
303                health_score: 1.0,
304                avg_latency_ms: 30.0,
305                ..ProviderStats::default()
306            },
307        ];
308
309        // Simulate health-based selection - choose highest health score
310        let best_provider = providers
311            .iter()
312            .enumerate()
313            .max_by(|(_, a), (_, b)| {
314                a.health_score
315                    .partial_cmp(&b.health_score)
316                    .unwrap_or(std::cmp::Ordering::Equal)
317            })
318            .map(|(idx, _)| idx);
319
320        assert_eq!(best_provider, Some(2)); // Third provider has 1.0 health score
321    }
322
323    #[test]
324    fn test_latency_based_selection_simulation() {
325        let providers = [
326            ProviderStats {
327                avg_latency_ms: 150.0,
328                ..ProviderStats::default()
329            },
330            ProviderStats {
331                avg_latency_ms: 50.0,
332                ..ProviderStats::default()
333            },
334            ProviderStats {
335                avg_latency_ms: 100.0,
336                ..ProviderStats::default()
337            },
338        ];
339
340        // Simulate least-latency selection
341        let fastest = providers
342            .iter()
343            .enumerate()
344            .min_by(|(_, a), (_, b)| {
345                a.avg_latency_ms
346                    .partial_cmp(&b.avg_latency_ms)
347                    .unwrap_or(std::cmp::Ordering::Equal)
348            })
349            .map(|(idx, _)| idx);
350
351        assert_eq!(fastest, Some(1)); // Second provider has lowest latency
352    }
353
354    #[test]
355    fn test_stats_aggregation() {
356        let providers = [
357            ProviderStats {
358                requests: 100,
359                errors: 5,
360                total_tokens: 50000,
361                total_cost: 5.0,
362                ..ProviderStats::default()
363            },
364            ProviderStats {
365                requests: 200,
366                errors: 10,
367                total_tokens: 100000,
368                total_cost: 10.0,
369                ..ProviderStats::default()
370            },
371        ];
372
373        let total_requests: u64 = providers.iter().map(|p| p.requests).sum();
374        let total_errors: u64 = providers.iter().map(|p| p.errors).sum();
375        let total_tokens: u64 = providers.iter().map(|p| p.total_tokens).sum();
376        let total_cost: f64 = providers.iter().map(|p| p.total_cost).sum();
377
378        assert_eq!(total_requests, 300);
379        assert_eq!(total_errors, 15);
380        assert_eq!(total_tokens, 150000);
381        assert!((total_cost - 15.0).abs() < f64::EPSILON);
382    }
383}
litellm_rs/sdk/client/types.rs

litellm_rs/sdk/client/
types.rs
