diff --git a/src/config/bot.rs b/src/config/bot.rs index 098c8d7..7177f92 100644 --- a/src/config/bot.rs +++ b/src/config/bot.rs @@ -129,6 +129,10 @@ impl RoleLike for Bot { &self.model } + fn model_mut(&mut self) -> &mut Model { + &mut self.model + } + fn temperature(&self) -> Option { self.config.temperature } diff --git a/src/config/mod.rs b/src/config/mod.rs index 7335774..a27245d 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -476,7 +476,7 @@ impl Config { match key { "max_output_tokens" => { let value = parse_value(value)?; - self.model.set_max_tokens(value, true); + self.set_max_output_tokens(value); } "temperature" => { let value = parse_value(value)?; @@ -564,6 +564,13 @@ impl Config { Ok(()) } + pub fn set_max_output_tokens(&mut self, value: Option) { + match self.role_like_mut() { + Some(role_like) => role_like.model_mut().set_max_tokens(value, true), + None => self.model.set_max_tokens(value, true), + }; + } + pub fn set_model(&mut self, model_id: &str) -> Result<()> { let model = Model::retrieve_chat(self, model_id)?; match self.role_like_mut() { diff --git a/src/config/role.rs b/src/config/role.rs index 4ff05b4..0dbfa00 100644 --- a/src/config/role.rs +++ b/src/config/role.rs @@ -18,6 +18,7 @@ pub const INPUT_PLACEHOLDER: &str = "__INPUT__"; pub trait RoleLike { fn to_role(&self) -> Role; fn model(&self) -> &Model; + fn model_mut(&mut self) -> &mut Model; fn temperature(&self) -> Option; fn top_p(&self) -> Option; fn functions_filter(&self) -> Option; @@ -229,6 +230,10 @@ impl RoleLike for Role { &self.model } + fn model_mut(&mut self) -> &mut Model { + &mut self.model + } + fn temperature(&self) -> Option { self.temperature } diff --git a/src/config/session.rs b/src/config/session.rs index df9b082..0b48cbd 100644 --- a/src/config/session.rs +++ b/src/config/session.rs @@ -426,6 +426,10 @@ impl RoleLike for Session { &self.model } + fn model_mut(&mut self) -> &mut Model { + &mut self.model + } + fn temperature(&self) -> Option { self.temperature }