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

1 files changed, 5 insertions, 0 deletions

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
     }