Skip to main content

litellm_rs/utils/data/
types.rs

1//! Type-safe wrappers and utilities
2//!
3//! This module provides type-safe wrappers for common values to prevent
4//! mixing up different types of IDs, keys, and other values.
5
6use serde::{Deserialize, Serialize};
7use std::fmt::{self, Display};
8use std::str::FromStr;
9use uuid::Uuid;
10
11/// A type-safe wrapper for user IDs
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
13pub struct UserId(Uuid);
14
15impl UserId {
16    /// Create a new user ID
17    pub fn new() -> Self {
18        Self(Uuid::new_v4())
19    }
20
21    /// Create a user ID from a UUID
22    pub fn from_uuid(uuid: Uuid) -> Self {
23        Self(uuid)
24    }
25
26    /// Get the inner UUID
27    pub fn as_uuid(&self) -> Uuid {
28        self.0
29    }
30
31    /// Get the string representation
32    pub fn as_str(&self) -> String {
33        self.0.to_string()
34    }
35}
36
37impl Default for UserId {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl Display for UserId {
44    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
45        write!(f, "{}", self.0)
46    }
47}
48
49impl FromStr for UserId {
50    type Err = uuid::Error;
51
52    fn from_str(s: &str) -> Result<Self, Self::Err> {
53        Ok(Self(Uuid::parse_str(s)?))
54    }
55}
56
57impl From<Uuid> for UserId {
58    fn from(uuid: Uuid) -> Self {
59        Self(uuid)
60    }
61}
62
63impl From<UserId> for Uuid {
64    fn from(user_id: UserId) -> Self {
65        user_id.0
66    }
67}
68
69/// A type-safe wrapper for API keys
70#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
71pub struct ApiKey(String);
72
73impl ApiKey {
74    /// Create a new API key from a string
75    pub fn new(key: String) -> Self {
76        Self(key)
77    }
78
79    /// Get the inner string
80    pub fn as_str(&self) -> &str {
81        &self.0
82    }
83
84    /// Get the string representation (for logging, truncated)
85    pub fn as_display_str(&self) -> String {
86        if self.0.len() > 8 {
87            format!("{}...", &self.0[..8])
88        } else {
89            self.0.clone()
90        }
91    }
92
93    /// Check if the key is valid (basic validation)
94    pub fn is_valid(&self) -> bool {
95        !self.0.is_empty() && self.0.len() >= 16
96    }
97}
98
99impl Display for ApiKey {
100    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
101        write!(f, "{}", self.as_display_str())
102    }
103}
104
105impl FromStr for ApiKey {
106    type Err = String;
107
108    fn from_str(s: &str) -> Result<Self, Self::Err> {
109        let key = Self(s.to_string());
110        if key.is_valid() {
111            Ok(key)
112        } else {
113            Err("Invalid API key format".to_string())
114        }
115    }
116}
117
118impl From<String> for ApiKey {
119    fn from(key: String) -> Self {
120        Self(key)
121    }
122}
123
124impl From<ApiKey> for String {
125    fn from(api_key: ApiKey) -> Self {
126        api_key.0
127    }
128}
129
130/// A type-safe wrapper for model names
131#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
132pub struct ModelName(String);
133
134impl ModelName {
135    /// Create a new model name
136    pub fn new(name: String) -> Self {
137        Self(name)
138    }
139
140    /// Get the inner string
141    pub fn as_str(&self) -> &str {
142        &self.0
143    }
144
145    /// Extract provider from model name (e.g., "openai/gpt-4" -> "openai")
146    pub fn provider(&self) -> Option<&str> {
147        self.0.find('/').map(|pos| &self.0[..pos])
148    }
149
150    /// Extract model from model name (e.g., "openai/gpt-4" -> "gpt-4")
151    pub fn model(&self) -> &str {
152        if let Some(pos) = self.0.find('/') {
153            &self.0[pos + 1..]
154        } else {
155            &self.0
156        }
157    }
158
159    /// Check if this is a valid model name
160    pub fn is_valid(&self) -> bool {
161        !self.0.is_empty() && !self.0.contains(' ')
162    }
163}
164
165impl Display for ModelName {
166    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
167        write!(f, "{}", self.0)
168    }
169}
170
171impl FromStr for ModelName {
172    type Err = String;
173
174    fn from_str(s: &str) -> Result<Self, Self::Err> {
175        let model = Self(s.to_string());
176        if model.is_valid() {
177            Ok(model)
178        } else {
179            Err("Invalid model name format".to_string())
180        }
181    }
182}
183
184impl From<String> for ModelName {
185    fn from(name: String) -> Self {
186        Self(name)
187    }
188}
189
190impl From<ModelName> for String {
191    fn from(model_name: ModelName) -> Self {
192        model_name.0
193    }
194}
195
196/// A type-safe wrapper for request IDs
197#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
198pub struct RequestId(String);
199
200impl RequestId {
201    /// Create a new request ID
202    pub fn new() -> Self {
203        Self(Uuid::new_v4().to_string())
204    }
205
206    /// Create a request ID from a string
207    pub fn from_string(id: String) -> Self {
208        Self(id)
209    }
210
211    /// Get the inner string
212    pub fn as_str(&self) -> &str {
213        &self.0
214    }
215}
216
217impl Default for RequestId {
218    fn default() -> Self {
219        Self::new()
220    }
221}
222
223impl Display for RequestId {
224    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
225        write!(f, "{}", self.0)
226    }
227}
228
229impl From<String> for RequestId {
230    fn from(id: String) -> Self {
231        Self(id)
232    }
233}
234
235impl From<RequestId> for String {
236    fn from(request_id: RequestId) -> Self {
237        request_id.0
238    }
239}
240
241/// A type-safe wrapper for organization IDs
242#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
243pub struct OrganizationId(Uuid);
244
245impl OrganizationId {
246    /// Create a new organization ID
247    pub fn new() -> Self {
248        Self(Uuid::new_v4())
249    }
250
251    /// Create an organization ID from a UUID
252    pub fn from_uuid(uuid: Uuid) -> Self {
253        Self(uuid)
254    }
255
256    /// Get the inner UUID
257    pub fn as_uuid(&self) -> Uuid {
258        self.0
259    }
260
261    /// Get the string representation
262    pub fn as_str(&self) -> String {
263        self.0.to_string()
264    }
265}
266
267impl Default for OrganizationId {
268    fn default() -> Self {
269        Self::new()
270    }
271}
272
273impl Display for OrganizationId {
274    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
275        write!(f, "{}", self.0)
276    }
277}
278
279impl FromStr for OrganizationId {
280    type Err = uuid::Error;
281
282    fn from_str(s: &str) -> Result<Self, Self::Err> {
283        Ok(Self(Uuid::parse_str(s)?))
284    }
285}
286
287impl From<Uuid> for OrganizationId {
288    fn from(uuid: Uuid) -> Self {
289        Self(uuid)
290    }
291}
292
293impl From<OrganizationId> for Uuid {
294    fn from(org_id: OrganizationId) -> Self {
295        org_id.0
296    }
297}
298
299/// A type-safe wrapper for team IDs
300#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
301pub struct TeamId(Uuid);
302
303impl TeamId {
304    /// Create a new team ID
305    pub fn new() -> Self {
306        Self(Uuid::new_v4())
307    }
308
309    /// Create a team ID from a UUID
310    pub fn from_uuid(uuid: Uuid) -> Self {
311        Self(uuid)
312    }
313
314    /// Get the inner UUID
315    pub fn as_uuid(&self) -> Uuid {
316        self.0
317    }
318
319    /// Get the string representation
320    pub fn as_str(&self) -> String {
321        self.0.to_string()
322    }
323}
324
325impl Default for TeamId {
326    fn default() -> Self {
327        Self::new()
328    }
329}
330
331impl Display for TeamId {
332    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
333        write!(f, "{}", self.0)
334    }
335}
336
337impl FromStr for TeamId {
338    type Err = uuid::Error;
339
340    fn from_str(s: &str) -> Result<Self, Self::Err> {
341        Ok(Self(Uuid::parse_str(s)?))
342    }
343}
344
345impl From<Uuid> for TeamId {
346    fn from(uuid: Uuid) -> Self {
347        Self(uuid)
348    }
349}
350
351impl From<TeamId> for Uuid {
352    fn from(team_id: TeamId) -> Self {
353        team_id.0
354    }
355}
356
357#[cfg(test)]
358mod tests {
359    use super::*;
360
361    #[test]
362    fn test_user_id() {
363        let user_id = UserId::new();
364        let uuid = user_id.as_uuid();
365        let user_id2 = UserId::from_uuid(uuid);
366        assert_eq!(user_id, user_id2);
367
368        let user_id_str = user_id.to_string();
369        let user_id3: UserId = user_id_str.parse().unwrap();
370        assert_eq!(user_id, user_id3);
371    }
372
373    #[test]
374    fn test_api_key() {
375        let key = ApiKey::new("sk-1234567890abcdef".to_string());
376        assert!(key.is_valid());
377        assert_eq!(key.as_str(), "sk-1234567890abcdef");
378        assert_eq!(key.as_display_str(), "sk-12345...");
379
380        let invalid_key = ApiKey::new("short".to_string());
381        assert!(!invalid_key.is_valid());
382    }
383
384    #[test]
385    fn test_model_name() {
386        let model = ModelName::new("openai/gpt-4".to_string());
387        assert!(model.is_valid());
388        assert_eq!(model.provider(), Some("openai"));
389        assert_eq!(model.model(), "gpt-4");
390
391        let simple_model = ModelName::new("gpt-4".to_string());
392        assert_eq!(simple_model.provider(), None);
393        assert_eq!(simple_model.model(), "gpt-4");
394    }
395
396    #[test]
397    fn test_request_id() {
398        let req_id = RequestId::new();
399        assert!(!req_id.as_str().is_empty());
400
401        let custom_id = RequestId::from_string("custom-123".to_string());
402        assert_eq!(custom_id.as_str(), "custom-123");
403    }
404}