refactor: make `set max_output_tokens` work for role/session/bot (#609)

1 files changed, 8 insertions, 1 deletions

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() {