bamboo_agent/core/
model_mapping.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3
4/// Anthropic model mapping configuration.
5///
6/// Used to map OpenAI-compatible model ids (e.g. "claude-3-opus") to the actual
7/// upstream Anthropic model id that should be used by the provider.
8#[derive(Debug, Clone, Serialize, Deserialize, Default)]
9pub struct AnthropicModelMapping {
10    #[serde(default)]
11    pub mappings: HashMap<String, String>,
12}
13
14/// Gemini model mapping configuration.
15///
16/// Used to map OpenAI-compatible model ids (e.g. "gemini-pro") to the actual
17/// upstream Gemini model id that should be used by the provider.
18#[derive(Debug, Clone, Serialize, Deserialize, Default)]
19pub struct GeminiModelMapping {
20    #[serde(default)]
21    pub mappings: HashMap<String, String>,
22}
23
24#[cfg(test)]
25mod tests {
26    use super::*;
27
28    #[test]
29    fn test_anthropic_model_mapping_default() {
30        let mapping = AnthropicModelMapping::default();
31        assert!(mapping.mappings.is_empty());
32    }
33
34    #[test]
35    fn test_gemini_model_mapping_default() {
36        let mapping = GeminiModelMapping::default();
37        assert!(mapping.mappings.is_empty());
38    }
39
40    #[test]
41    fn test_anthropic_model_mapping_with_entries() {
42        let mut mappings = HashMap::new();
43        mappings.insert(
44            "claude-3-opus".to_string(),
45            "claude-3-opus-20240229".to_string(),
46        );
47        mappings.insert(
48            "claude-3-sonnet".to_string(),
49            "claude-3-sonnet-20240229".to_string(),
50        );
51
52        let mapping = AnthropicModelMapping { mappings };
53        assert_eq!(mapping.mappings.len(), 2);
54        assert_eq!(
55            mapping.mappings.get("claude-3-opus"),
56            Some(&"claude-3-opus-20240229".to_string())
57        );
58    }
59
60    #[test]
61    fn test_gemini_model_mapping_with_entries() {
62        let mut mappings = HashMap::new();
63        mappings.insert("gemini-pro".to_string(), "gemini-1.5-pro".to_string());
64        mappings.insert("gemini-flash".to_string(), "gemini-1.5-flash".to_string());
65
66        let mapping = GeminiModelMapping { mappings };
67        assert_eq!(mapping.mappings.len(), 2);
68        assert_eq!(
69            mapping.mappings.get("gemini-pro"),
70            Some(&"gemini-1.5-pro".to_string())
71        );
72    }
73
74    #[test]
75    fn test_anthropic_model_mapping_serialization() {
76        let mut mappings = HashMap::new();
77        mappings.insert("test-model".to_string(), "actual-model".to_string());
78
79        let mapping = AnthropicModelMapping { mappings };
80        let json = serde_json::to_string(&mapping).unwrap();
81
82        assert!(json.contains("test-model"));
83        assert!(json.contains("actual-model"));
84    }
85
86    #[test]
87    fn test_gemini_model_mapping_serialization() {
88        let mut mappings = HashMap::new();
89        mappings.insert("gemini-test".to_string(), "gemini-actual".to_string());
90
91        let mapping = GeminiModelMapping { mappings };
92        let json = serde_json::to_string(&mapping).unwrap();
93
94        assert!(json.contains("gemini-test"));
95        assert!(json.contains("gemini-actual"));
96    }
97
98    #[test]
99    fn test_anthropic_model_mapping_deserialization() {
100        let json = r#"{"mappings":{"key":"value"}}"#;
101        let mapping: AnthropicModelMapping = serde_json::from_str(json).unwrap();
102
103        assert_eq!(mapping.mappings.len(), 1);
104        assert_eq!(mapping.mappings.get("key"), Some(&"value".to_string()));
105    }
106
107    #[test]
108    fn test_gemini_model_mapping_deserialization() {
109        let json = r#"{"mappings":{"g-key":"g-value"}}"#;
110        let mapping: GeminiModelMapping = serde_json::from_str(json).unwrap();
111
112        assert_eq!(mapping.mappings.len(), 1);
113        assert_eq!(mapping.mappings.get("g-key"), Some(&"g-value".to_string()));
114    }
115
116    #[test]
117    fn test_anthropic_model_mapping_empty_json() {
118        let json = r#"{}"#;
119        let mapping: AnthropicModelMapping = serde_json::from_str(json).unwrap();
120
121        assert!(mapping.mappings.is_empty());
122    }
123
124    #[test]
125    fn test_gemini_model_mapping_empty_json() {
126        let json = r#"{}"#;
127        let mapping: GeminiModelMapping = serde_json::from_str(json).unwrap();
128
129        assert!(mapping.mappings.is_empty());
130    }
131
132    #[test]
133    fn test_anthropic_model_mapping_clone() {
134        let mut mappings = HashMap::new();
135        mappings.insert("key".to_string(), "value".to_string());
136
137        let mapping = AnthropicModelMapping { mappings };
138        let cloned = mapping.clone();
139
140        assert_eq!(mapping.mappings, cloned.mappings);
141    }
142
143    #[test]
144    fn test_gemini_model_mapping_clone() {
145        let mut mappings = HashMap::new();
146        mappings.insert("gkey".to_string(), "gvalue".to_string());
147
148        let mapping = GeminiModelMapping { mappings };
149        let cloned = mapping.clone();
150
151        assert_eq!(mapping.mappings, cloned.mappings);
152    }
153
154    #[test]
155    fn test_anthropic_model_mapping_debug() {
156        let mapping = AnthropicModelMapping::default();
157        let debug_str = format!("{:?}", mapping);
158
159        assert!(debug_str.contains("AnthropicModelMapping"));
160    }
161
162    #[test]
163    fn test_gemini_model_mapping_debug() {
164        let mapping = GeminiModelMapping::default();
165        let debug_str = format!("{:?}", mapping);
166
167        assert!(debug_str.contains("GeminiModelMapping"));
168    }
169
170    #[test]
171    fn test_anthropic_model_mapping_multiple_entries() {
172        let mut mappings = HashMap::new();
173        for i in 0..10 {
174            mappings.insert(format!("model-{}", i), format!("actual-model-{}", i));
175        }
176
177        let mapping = AnthropicModelMapping { mappings };
178        assert_eq!(mapping.mappings.len(), 10);
179    }
180
181    #[test]
182    fn test_gemini_model_mapping_multiple_entries() {
183        let mut mappings = HashMap::new();
184        for i in 0..10 {
185            mappings.insert(format!("gemini-{}", i), format!("gemini-actual-{}", i));
186        }
187
188        let mapping = GeminiModelMapping { mappings };
189        assert_eq!(mapping.mappings.len(), 10);
190    }
191}
bamboo_agent/core/model_mapping.rs

bamboo_agent/core/
model_mapping.rs
