mirror of
https://wget.la/https://github.com/leookun/cursor-byok
synced 2026-10-04 19:31:28 +08:00
refactor: share token count utilities
This commit is contained in:
@@ -14,7 +14,7 @@ use crate::{
|
||||
proxy::{self, CursorProxy},
|
||||
CursorSessionRegistry,
|
||||
},
|
||||
model::ProviderModel,
|
||||
model::{format_token_count, parse_token_count, ProviderModel},
|
||||
Error, Result,
|
||||
};
|
||||
|
||||
@@ -214,7 +214,7 @@ fn context_options(model: &ProviderModel) -> Vec<(String, String)> {
|
||||
let value = tokens.to_string();
|
||||
let duplicate = contexts
|
||||
.iter()
|
||||
.any(|(existing, _)| parse_context_tokens(existing) == Some(tokens));
|
||||
.any(|(existing, _)| parse_token_count(existing) == Some(tokens));
|
||||
if !duplicate {
|
||||
contexts.push((value, format!("{} (Custom)", format_token_count(tokens))));
|
||||
}
|
||||
@@ -222,26 +222,6 @@ fn context_options(model: &ProviderModel) -> Vec<(String, String)> {
|
||||
contexts
|
||||
}
|
||||
|
||||
fn parse_context_tokens(value: &str) -> Option<u64> {
|
||||
let value = value.trim().to_ascii_lowercase();
|
||||
let (number, multiplier) = match value.chars().last()? {
|
||||
'k' => (&value[..value.len() - 1], 1_000),
|
||||
'm' => (&value[..value.len() - 1], 1_000_000),
|
||||
_ => (value.as_str(), 1),
|
||||
};
|
||||
number.parse::<u64>().ok()?.checked_mul(multiplier)
|
||||
}
|
||||
|
||||
fn format_token_count(tokens: u64) -> String {
|
||||
if tokens >= 1_000_000 && tokens.is_multiple_of(1_000_000) {
|
||||
format!("{}M", tokens / 1_000_000)
|
||||
} else if tokens >= 1_000 && tokens.is_multiple_of(1_000) {
|
||||
format!("{}K", tokens / 1_000)
|
||||
} else {
|
||||
tokens.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn available_models(
|
||||
State(registry): State<CursorSessionRegistry>,
|
||||
Extension(proxy): Extension<CursorProxy>,
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
use crate::{
|
||||
cursor::proto::agent::v1 as pb,
|
||||
model::{ModelLatency, ModelSpec, ReasoningSpec, SubagentKind, SubagentModelOverride},
|
||||
model::{
|
||||
parse_token_count, ModelLatency, ModelSpec, ReasoningSpec, SubagentKind,
|
||||
SubagentModelOverride,
|
||||
},
|
||||
Error, Result,
|
||||
};
|
||||
|
||||
@@ -136,16 +139,6 @@ fn parse_bool(parameter: &pb::requested_model::ModelParameterValue) -> Result<bo
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_token_count(value: &str) -> Option<u64> {
|
||||
let value = value.trim().to_ascii_lowercase();
|
||||
let (number, multiplier) = match value.chars().last()? {
|
||||
'k' => (&value[..value.len() - 1], 1_000),
|
||||
'm' => (&value[..value.len() - 1], 1_000_000),
|
||||
_ => (value.as_str(), 1),
|
||||
};
|
||||
number.parse::<u64>().ok()?.checked_mul(multiplier)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
@@ -9,6 +9,7 @@ mod projection;
|
||||
mod provider;
|
||||
mod run;
|
||||
mod runtime_tag;
|
||||
mod token_count;
|
||||
mod tool;
|
||||
mod usage;
|
||||
|
||||
@@ -23,5 +24,6 @@ pub use projection::*;
|
||||
pub use provider::*;
|
||||
pub use run::*;
|
||||
pub use runtime_tag::*;
|
||||
pub(crate) use token_count::*;
|
||||
pub use tool::*;
|
||||
pub use usage::*;
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
pub(crate) fn parse_token_count(value: &str) -> Option<u64> {
|
||||
let value = value.trim().to_ascii_lowercase();
|
||||
let (number, multiplier) = match value.chars().last()? {
|
||||
'k' => (&value[..value.len() - 1], 1_000),
|
||||
'm' => (&value[..value.len() - 1], 1_000_000),
|
||||
_ => (value.as_str(), 1),
|
||||
};
|
||||
number.parse::<u64>().ok()?.checked_mul(multiplier)
|
||||
}
|
||||
|
||||
pub(crate) fn format_token_count(tokens: u64) -> String {
|
||||
if tokens >= 1_000_000 && tokens.is_multiple_of(1_000_000) {
|
||||
format!("{}M", tokens / 1_000_000)
|
||||
} else if tokens >= 1_000 && tokens.is_multiple_of(1_000) {
|
||||
format!("{}K", tokens / 1_000)
|
||||
} else {
|
||||
tokens.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn token_counts_parse_plain_and_abbreviated_values() {
|
||||
assert_eq!(parse_token_count("272000"), Some(272_000));
|
||||
assert_eq!(parse_token_count("272K"), Some(272_000));
|
||||
assert_eq!(parse_token_count("1m"), Some(1_000_000));
|
||||
assert_eq!(parse_token_count("invalid"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_counts_format_exact_thousands_and_millions() {
|
||||
assert_eq!(format_token_count(272_000), "272K");
|
||||
assert_eq!(format_token_count(1_000_000), "1M");
|
||||
assert_eq!(format_token_count(272_001), "272001");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user