litellm_rs/core/models/
mod.rs

1//! Core data models for the Gateway
2//!
3//! This module defines all the core data structures used throughout the gateway.
4
5pub mod deployment;
6pub mod metrics;
7pub mod openai;
8pub mod request;
9pub mod response;
10pub mod team;
11pub mod user;
12
13// Re-export commonly used types
14
15pub use team::*;
16pub use user::*;
17
18use serde::{Deserialize, Serialize};
19use std::collections::HashMap;
20use uuid::Uuid;
21
22/// Common metadata for all models
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct Metadata {
25    /// Unique identifier
26    pub id: Uuid,
27    /// Creation timestamp
28    pub created_at: chrono::DateTime<chrono::Utc>,
29    /// Last update timestamp
30    pub updated_at: chrono::DateTime<chrono::Utc>,
31    /// Version for optimistic locking
32    pub version: i64,
33    /// Additional metadata
34    #[serde(default)]
35    pub extra: HashMap<String, serde_json::Value>,
36}
37
38impl Default for Metadata {
39    fn default() -> Self {
40        let now = chrono::Utc::now();
41        Self {
42            id: Uuid::new_v4(),
43            created_at: now,
44            updated_at: now,
45            version: 1,
46            extra: HashMap::new(),
47        }
48    }
49}
50
51impl Metadata {
52    /// Create new metadata
53    pub fn new() -> Self {
54        Self::default()
55    }
56
57    /// Update the timestamp and increment version
58    pub fn touch(&mut self) {
59        self.updated_at = chrono::Utc::now();
60        self.version += 1;
61    }
62
63    /// Set extra metadata
64    pub fn set_extra<K: Into<String>, V: Into<serde_json::Value>>(&mut self, key: K, value: V) {
65        self.extra.insert(key.into(), value.into());
66    }
67
68    /// Get extra metadata
69    pub fn get_extra(&self, key: &str) -> Option<&serde_json::Value> {
70        self.extra.get(key)
71    }
72}
73
74/// API Key information
75#[derive(Debug, Clone, Serialize, Deserialize)]
76pub struct ApiKey {
77    /// Metadata
78    #[serde(flatten)]
79    pub metadata: Metadata,
80    /// Key name/description
81    pub name: String,
82    /// Hashed key value
83    pub key_hash: String,
84    /// Key prefix for identification
85    pub key_prefix: String,
86    /// Associated user ID
87    pub user_id: Option<Uuid>,
88    /// Associated team ID
89    pub team_id: Option<Uuid>,
90    /// Permissions
91    pub permissions: Vec<String>,
92    /// Rate limits
93    pub rate_limits: Option<RateLimits>,
94    /// Expiration date
95    pub expires_at: Option<chrono::DateTime<chrono::Utc>>,
96    /// Whether the key is active
97    pub is_active: bool,
98    /// Last used timestamp
99    pub last_used_at: Option<chrono::DateTime<chrono::Utc>>,
100    /// Usage statistics
101    pub usage_stats: UsageStats,
102}
103
104/// Rate limits for API keys
105#[derive(Debug, Clone, Serialize, Deserialize)]
106pub struct RateLimits {
107    /// Requests per minute
108    pub rpm: Option<u32>,
109    /// Tokens per minute
110    pub tpm: Option<u32>,
111    /// Requests per day
112    pub rpd: Option<u32>,
113    /// Tokens per day
114    pub tpd: Option<u32>,
115    /// Concurrent requests
116    pub concurrent: Option<u32>,
117}
118
119/// Usage statistics
120#[derive(Debug, Clone, Serialize, Deserialize, Default)]
121pub struct UsageStats {
122    /// Total requests
123    pub total_requests: u64,
124    /// Total tokens
125    pub total_tokens: u64,
126    /// Total cost
127    pub total_cost: f64,
128    /// Requests today
129    pub requests_today: u32,
130    /// Tokens today
131    pub tokens_today: u32,
132    /// Cost today
133    pub cost_today: f64,
134    /// Last reset date
135    pub last_reset: chrono::DateTime<chrono::Utc>,
136}
137
138/// Model information
139#[derive(Debug, Clone, Serialize, Deserialize)]
140pub struct ModelInfo {
141    /// Model ID
142    pub id: String,
143    /// Model name
144    pub name: String,
145    /// Provider
146    pub provider: String,
147    /// Model type (chat, completion, embedding, etc.)
148    pub model_type: ModelType,
149    /// Context window size
150    pub context_window: Option<u32>,
151    /// Maximum output tokens
152    pub max_output_tokens: Option<u32>,
153    /// Input cost per token
154    pub input_cost_per_token: Option<f64>,
155    /// Output cost per token
156    pub output_cost_per_token: Option<f64>,
157    /// Supported features
158    pub features: Vec<String>,
159    /// Model description
160    pub description: Option<String>,
161    /// Whether the model is available
162    pub is_available: bool,
163}
164
165/// Model types
166#[derive(Debug, Clone, Serialize, Deserialize)]
167#[serde(rename_all = "snake_case")]
168pub enum ModelType {
169    /// Chat completion models
170    Chat,
171    /// Text completion models
172    Completion,
173    /// Text embedding models
174    Embedding,
175    /// Image generation models
176    ImageGeneration,
177    /// Audio transcription models
178    AudioTranscription,
179    /// Audio translation models
180    AudioTranslation,
181    /// Content moderation models
182    Moderation,
183    /// Fine-tuning models
184    FineTuning,
185    /// Document reranking models
186    Rerank,
187}
188
189/// Request context information
190#[derive(Debug, Clone, Serialize, Deserialize)]
191pub struct RequestContext {
192    /// Request ID
193    pub request_id: String,
194    /// User ID (if authenticated)
195    pub user_id: Option<Uuid>,
196    /// Team ID (if applicable)
197    pub team_id: Option<Uuid>,
198    /// API Key ID (if using API key auth)
199    pub api_key_id: Option<Uuid>,
200    /// Client IP address
201    pub client_ip: Option<String>,
202    /// User agent
203    pub user_agent: Option<String>,
204    /// Request timestamp
205    pub timestamp: chrono::DateTime<chrono::Utc>,
206    /// Additional headers
207    pub headers: HashMap<String, String>,
208    /// Trace ID for distributed tracing
209    pub trace_id: Option<String>,
210    /// Span ID for distributed tracing
211    pub span_id: Option<String>,
212}
213
214impl Default for RequestContext {
215    fn default() -> Self {
216        Self {
217            request_id: Uuid::new_v4().to_string(),
218            user_id: None,
219            team_id: None,
220            api_key_id: None,
221            client_ip: None,
222            user_agent: None,
223            timestamp: chrono::Utc::now(),
224            headers: HashMap::new(),
225            trace_id: None,
226            span_id: None,
227        }
228    }
229}
230
231impl RequestContext {
232    /// Create new request context
233    pub fn new() -> Self {
234        Self::default()
235    }
236
237    /// Set user information
238    pub fn with_user(mut self, user_id: Uuid, team_id: Option<Uuid>) -> Self {
239        self.user_id = Some(user_id);
240        self.team_id = team_id;
241        self
242    }
243
244    /// Set API key information
245    pub fn with_api_key(mut self, api_key_id: Uuid) -> Self {
246        self.api_key_id = Some(api_key_id);
247        self
248    }
249
250    /// Set client information
251    pub fn with_client_info(mut self, ip: Option<String>, user_agent: Option<String>) -> Self {
252        self.client_ip = ip;
253        self.user_agent = user_agent;
254        self
255    }
256
257    /// Set tracing information
258    pub fn with_tracing(mut self, trace_id: String, span_id: String) -> Self {
259        self.trace_id = Some(trace_id);
260        self.span_id = Some(span_id);
261        self
262    }
263
264    /// Add header
265    pub fn add_header<K: Into<String>, V: Into<String>>(mut self, key: K, value: V) -> Self {
266        self.headers.insert(key.into(), value.into());
267        self
268    }
269}
270
271/// Provider health information
272#[derive(Debug, Clone, Serialize, Deserialize)]
273pub struct ProviderHealth {
274    /// Provider name
275    pub provider: String,
276    /// Health status
277    pub status: HealthStatus,
278    /// Last check timestamp
279    pub last_check: chrono::DateTime<chrono::Utc>,
280    /// Response time in milliseconds
281    pub response_time_ms: Option<u64>,
282    /// Error message if unhealthy
283    pub error_message: Option<String>,
284    /// Success rate (0.0 to 1.0)
285    pub success_rate: f64,
286    /// Total requests in the last period
287    pub total_requests: u64,
288    /// Failed requests in the last period
289    pub failed_requests: u64,
290}
291
292/// Health status enum
293#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
294#[serde(rename_all = "lowercase")]
295pub enum HealthStatus {
296    /// Service is operating normally
297    Healthy,
298    /// Service is experiencing minor issues
299    Degraded,
300    /// Service is not functioning properly
301    Unhealthy,
302    /// Health status cannot be determined
303    Unknown,
304}
305
306/// Provider registry health summary
307#[derive(Debug, Clone, Serialize, Deserialize)]
308pub struct ProviderRegistryHealth {
309    /// Total number of providers
310    pub total_count: usize,
311    /// Number of healthy providers
312    pub healthy_count: usize,
313    /// Number of degraded providers
314    pub degraded_count: usize,
315    /// Number of unhealthy providers
316    pub unhealthy_count: usize,
317    /// Individual provider health
318    pub providers: Vec<ProviderHealth>,
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324
325    #[test]
326    fn test_metadata_creation() {
327        let metadata = Metadata::new();
328        assert_eq!(metadata.version, 1);
329        assert!(metadata.created_at <= chrono::Utc::now());
330    }
331
332    #[test]
333    fn test_metadata_touch() {
334        let mut metadata = Metadata::new();
335        let original_version = metadata.version;
336        let original_updated = metadata.updated_at;
337
338        std::thread::sleep(std::time::Duration::from_millis(1));
339        metadata.touch();
340
341        assert_eq!(metadata.version, original_version + 1);
342        assert!(metadata.updated_at > original_updated);
343    }
344
345    #[test]
346    fn test_request_context_creation() {
347        let context = RequestContext::new();
348        assert!(!context.request_id.is_empty());
349        assert!(context.user_id.is_none());
350    }
351
352    #[test]
353    fn test_request_context_builder() {
354        let user_id = Uuid::new_v4();
355        let team_id = Uuid::new_v4();
356
357        let context = RequestContext::new()
358            .with_user(user_id, Some(team_id))
359            .with_client_info(
360                Some("127.0.0.1".to_string()),
361                Some("test-agent".to_string()),
362            )
363            .add_header("X-Custom", "value");
364
365        assert_eq!(context.user_id, Some(user_id));
366        assert_eq!(context.team_id, Some(team_id));
367        assert_eq!(context.client_ip, Some("127.0.0.1".to_string()));
368        assert_eq!(context.headers.get("X-Custom"), Some(&"value".to_string()));
369    }
370}
litellm_rs/core/models/mod.rs

litellm_rs/core/models/
mod.rs
