Skip to content

Commit

Permalink
refactor: make set max_output_tokens work for role/session/bot (#609)
Browse files Browse the repository at this point in the history
  • Loading branch information
sigoden committed Jun 17, 2024
1 parent ba884c9 commit ff28477
Show file tree
Hide file tree
Showing 4 changed files with 21 additions and 1 deletion.
4 changes: 4 additions & 0 deletions src/config/bot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
9 changes: 8 additions & 1 deletion src/config/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)?;
Expand Down Expand Up @@ -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() {
Expand Down
5 changes: 5 additions & 0 deletions src/config/role.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>;
Expand Down Expand Up @@ -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
}
Expand Down
4 changes: 4 additions & 0 deletions src/config/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down

0 comments on commit ff28477

Please sign in to comment.