4 files changed, 21 insertions, 1 deletions
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<f64> {
         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<isize>) {
+        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<f64>;
     fn top_p(&self) -> Option<f64>;
     fn functions_filter(&self) -> Option<FunctionsFilter>;
@@ -229,6 +230,10 @@ impl RoleLike for Role {
         &self.model
     }
 
+    fn model_mut(&mut self) -> &mut Model {
+        &mut self.model
+    }
+
     fn temperature(&self) -> Option<f64> {
         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<f64> {
         self.temperature
     }