diff options
Diffstat (limited to 'src/client/model.rs')
| -rw-r--r-- | src/client/model.rs | 14 |
1 files changed, 11 insertions, 3 deletions
diff --git a/src/client/model.rs b/src/client/model.rs index b562705..bcd474f 100644 --- a/src/client/model.rs +++ b/src/client/model.rs @@ -5,7 +5,7 @@ use super::{ }; use crate::config::Config; -use crate::utils::estimate_token_length; +use crate::utils::{estimate_token_length, strip_think_tag}; use anyhow::{bail, Result}; use serde::{Deserialize, Serialize}; @@ -220,10 +220,18 @@ impl Model { } pub fn messages_tokens(&self, messages: &[Message]) -> usize { + let messages_len = messages.len(); messages .iter() - .map(|v| match &v.content { - MessageContent::Text(text) => estimate_token_length(text), + .enumerate() + .map(|(i, v)| match &v.content { + MessageContent::Text(text) => { + if v.role.is_assistant() && i != messages_len - 1 { + estimate_token_length(&strip_think_tag(text)) + } else { + estimate_token_length(text) + } + } MessageContent::Array(list) => list .iter() .map(|v| match v { |