orchestra_rs/model/
config.rs

1use serde::{Deserialize, Serialize};
2use crate::error::{OrchestraError, Result};
3
4/// Configuration for a language model
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub struct ModelConfig {
7    pub name: String,
8    pub system_instruction: Option<String>,
9    pub temperature: f32,
10    pub top_p: f32,
11    pub top_k: Option<u32>,
12    pub max_tokens: Option<u32>,
13    pub thinking_mode: Option<bool>,
14    pub stop_sequences: Vec<String>,
15}
16
17impl Default for ModelConfig {
18    fn default() -> Self {
19        ModelConfig {
20            name: String::new(),
21            system_instruction: None,
22            temperature: 1.0,
23            top_p: 0.95,
24            top_k: None,
25            max_tokens: None,
26            thinking_mode: None,
27            stop_sequences: Vec::new(),
28        }
29    }
30}
31
32impl ModelConfig {
33    /// Create a new model configuration with the given model name
34    pub fn new<S: Into<String>>(name: S) -> Self {
35        Self {
36            name: name.into(),
37            ..Default::default()
38        }
39    }
40
41    /// Set the model name
42    pub fn with_name<S: Into<String>>(mut self, name: S) -> Self {
43        self.name = name.into();
44        self
45    }
46
47    /// Set the system instruction
48    pub fn with_system_instruction<S: Into<String>>(mut self, instruction: S) -> Self {
49        self.system_instruction = Some(instruction.into());
50        self
51    }
52
53    /// Set the temperature (0.0 to 2.0)
54    pub fn with_temperature(mut self, temperature: f32) -> Result<Self> {
55        if !(0.0..=2.0).contains(&temperature) {
56            return Err(OrchestraError::config("Temperature must be between 0.0 and 2.0"));
57        }
58        self.temperature = temperature;
59        Ok(self)
60    }
61
62    /// Set the top_p (0.0 to 1.0)
63    pub fn with_top_p(mut self, top_p: f32) -> Result<Self> {
64        if !(0.0..=1.0).contains(&top_p) {
65            return Err(OrchestraError::config("top_p must be between 0.0 and 1.0"));
66        }
67        self.top_p = top_p;
68        Ok(self)
69    }
70
71    /// Set the top_k
72    pub fn with_top_k(mut self, top_k: u32) -> Self {
73        self.top_k = Some(top_k);
74        self
75    }
76
77    /// Set the maximum number of tokens to generate
78    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
79        self.max_tokens = Some(max_tokens);
80        self
81    }
82
83    /// Enable or disable thinking mode
84    pub fn with_thinking_mode(mut self, thinking_mode: bool) -> Self {
85        self.thinking_mode = Some(thinking_mode);
86        self
87    }
88
89    /// Add a stop sequence
90    pub fn with_stop_sequence<S: Into<String>>(mut self, stop_sequence: S) -> Self {
91        self.stop_sequences.push(stop_sequence.into());
92        self
93    }
94
95    /// Set multiple stop sequences
96    pub fn with_stop_sequences<I, S>(mut self, stop_sequences: I) -> Self
97    where
98        I: IntoIterator<Item = S>,
99        S: Into<String>,
100    {
101        self.stop_sequences = stop_sequences.into_iter().map(|s| s.into()).collect();
102        self
103    }
104
105    /// Validate the configuration
106    pub fn validate(&self) -> Result<()> {
107        if self.name.is_empty() {
108            return Err(OrchestraError::config("Model name cannot be empty"));
109        }
110
111        if !(0.0..=2.0).contains(&self.temperature) {
112            return Err(OrchestraError::config("Temperature must be between 0.0 and 2.0"));
113        }
114
115        if !(0.0..=1.0).contains(&self.top_p) {
116            return Err(OrchestraError::config("top_p must be between 0.0 and 1.0"));
117        }
118
119        if let Some(max_tokens) = self.max_tokens {
120            if max_tokens == 0 {
121                return Err(OrchestraError::config("max_tokens must be greater than 0"));
122            }
123        }
124
125        Ok(())
126    }
127
128    /// Create a conservative configuration (lower temperature, more focused)
129    pub fn conservative<S: Into<String>>(name: S) -> Self {
130        Self::new(name)
131            .with_temperature(0.3)
132            .unwrap()
133            .with_top_p(0.8)
134            .unwrap()
135    }
136
137    /// Create a creative configuration (higher temperature, more diverse)
138    pub fn creative<S: Into<String>>(name: S) -> Self {
139        Self::new(name)
140            .with_temperature(1.2)
141            .unwrap()
142            .with_top_p(0.95)
143            .unwrap()
144    }
145
146    /// Create a balanced configuration (moderate settings)
147    pub fn balanced<S: Into<String>>(name: S) -> Self {
148        Self::new(name)
149            .with_temperature(0.7)
150            .unwrap()
151            .with_top_p(0.9)
152            .unwrap()
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159
160    #[test]
161    fn test_model_config_new() {
162        let config = ModelConfig::new("test-model");
163        assert_eq!(config.name, "test-model");
164        assert_eq!(config.temperature, 1.0);
165        assert_eq!(config.top_p, 0.95);
166        assert!(config.system_instruction.is_none());
167        assert!(config.top_k.is_none());
168        assert!(config.max_tokens.is_none());
169        assert!(config.thinking_mode.is_none());
170        assert!(config.stop_sequences.is_empty());
171    }
172
173    #[test]
174    fn test_model_config_builder() {
175        let config = ModelConfig::new("test-model")
176            .with_system_instruction("You are helpful")
177            .with_temperature(0.5)
178            .unwrap()
179            .with_top_p(0.8)
180            .unwrap()
181            .with_top_k(40)
182            .with_max_tokens(1000)
183            .with_thinking_mode(true)
184            .with_stop_sequence("STOP");
185
186        assert_eq!(config.name, "test-model");
187        assert_eq!(config.system_instruction, Some("You are helpful".to_string()));
188        assert_eq!(config.temperature, 0.5);
189        assert_eq!(config.top_p, 0.8);
190        assert_eq!(config.top_k, Some(40));
191        assert_eq!(config.max_tokens, Some(1000));
192        assert_eq!(config.thinking_mode, Some(true));
193        assert_eq!(config.stop_sequences, vec!["STOP"]);
194    }
195
196    #[test]
197    fn test_model_config_validation() {
198        let config = ModelConfig::new("test-model");
199        assert!(config.validate().is_ok());
200
201        // Test empty name
202        let mut config = ModelConfig::new("");
203        assert!(config.validate().is_err());
204
205        // Test invalid temperature
206        config = ModelConfig::new("test");
207        config.temperature = 3.0;
208        assert!(config.validate().is_err());
209
210        config.temperature = -1.0;
211        assert!(config.validate().is_err());
212
213        // Test invalid top_p
214        config = ModelConfig::new("test");
215        config.top_p = 1.5;
216        assert!(config.validate().is_err());
217
218        config.top_p = -0.1;
219        assert!(config.validate().is_err());
220
221        // Test invalid max_tokens
222        config = ModelConfig::new("test");
223        config.max_tokens = Some(0);
224        assert!(config.validate().is_err());
225    }
226
227    #[test]
228    fn test_model_config_temperature_validation() {
229        let config = ModelConfig::new("test");
230
231        // Valid temperatures
232        assert!(config.clone().with_temperature(0.0).is_ok());
233        assert!(config.clone().with_temperature(1.0).is_ok());
234        assert!(config.clone().with_temperature(2.0).is_ok());
235
236        // Invalid temperatures
237        assert!(config.clone().with_temperature(-0.1).is_err());
238        assert!(config.clone().with_temperature(2.1).is_err());
239    }
240
241    #[test]
242    fn test_model_config_top_p_validation() {
243        let config = ModelConfig::new("test");
244
245        // Valid top_p values
246        assert!(config.clone().with_top_p(0.0).is_ok());
247        assert!(config.clone().with_top_p(0.5).is_ok());
248        assert!(config.clone().with_top_p(1.0).is_ok());
249
250        // Invalid top_p values
251        assert!(config.clone().with_top_p(-0.1).is_err());
252        assert!(config.clone().with_top_p(1.1).is_err());
253    }
254
255    #[test]
256    fn test_model_config_presets() {
257        let conservative = ModelConfig::conservative("test-model");
258        assert_eq!(conservative.temperature, 0.3);
259        assert_eq!(conservative.top_p, 0.8);
260
261        let creative = ModelConfig::creative("test-model");
262        assert_eq!(creative.temperature, 1.2);
263        assert_eq!(creative.top_p, 0.95);
264
265        let balanced = ModelConfig::balanced("test-model");
266        assert_eq!(balanced.temperature, 0.7);
267        assert_eq!(balanced.top_p, 0.9);
268    }
269
270    #[test]
271    fn test_model_config_stop_sequences() {
272        let config = ModelConfig::new("test")
273            .with_stop_sequences(vec!["STOP", "END", "FINISH"]);
274
275        assert_eq!(config.stop_sequences.len(), 3);
276        assert!(config.stop_sequences.contains(&"STOP".to_string()));
277        assert!(config.stop_sequences.contains(&"END".to_string()));
278        assert!(config.stop_sequences.contains(&"FINISH".to_string()));
279    }
280}
orchestra_rs/model/config.rs

orchestra_rs/model/
config.rs
