1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
|
use super::Message;
use crate::utils::count_tokens;
use anyhow::{bail, Result};
#[derive(Debug, Clone)]
pub struct ModelInfo {
pub client: String,
pub name: String,
pub index: usize,
pub max_tokens: Option<usize>,
pub per_message_tokens: usize,
pub bias_tokens: usize,
}
impl Default for ModelInfo {
fn default() -> Self {
ModelInfo::new(0, "", "")
}
}
impl ModelInfo {
pub fn new(index: usize, client: &str, name: &str) -> Self {
Self {
index,
client: client.into(),
name: name.into(),
max_tokens: None,
per_message_tokens: 0,
bias_tokens: 0,
}
}
pub fn set_max_tokens(mut self, max_tokens: Option<usize>) -> Self {
match max_tokens {
None | Some(0) => self.max_tokens = None,
_ => self.max_tokens = max_tokens,
}
self
}
pub fn set_tokens_formula(mut self, per_message_token: usize, bias_tokens: usize) -> Self {
self.per_message_tokens = per_message_token;
self.bias_tokens = bias_tokens;
self
}
pub fn full_name(&self) -> String {
format!("{}:{}", self.client, self.name)
}
pub fn messages_tokens(&self, messages: &[Message]) -> usize {
messages.iter().map(|v| count_tokens(&v.content)).sum()
}
pub fn totatl_tokens(&self, messages: &[Message]) -> usize {
if messages.is_empty() {
return 0;
}
let num_messages = messages.len();
let message_tokens = self.messages_tokens(messages);
if messages[num_messages - 1].role.is_user() {
num_messages * self.per_message_tokens + message_tokens
} else {
(num_messages - 1) * self.per_message_tokens + message_tokens
}
}
pub fn max_tokens_limit(&self, messages: &[Message]) -> Result<()> {
let total_tokens = self.totatl_tokens(messages) + self.bias_tokens;
if let Some(max_tokens) = self.max_tokens {
if total_tokens >= max_tokens {
bail!("Exceed max tokens limit")
}
}
Ok(())
}
}
|
