From ff284779d9757c128acaf690ddf64a87a986e413 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Mon, 17 Jun 2024 15:17:07 +0800
Subject: refactor: make `set max_output_tokens` work for role/session/bot
 (#609)

---
 src/config/mod.rs | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

(limited to 'src/config/mod.rs')

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() {
-- 
cgit v1.2.3