collab: Clean up LLM token creation (#26955)
This PR cleans up the LLM token creation a bit. We now pass in the entire list of feature flags to the `LlmTokenClaims::create` method to prevent having a bunch of confusable `bool` parameters. Release Notes: - N/A
This commit is contained in:
parent
2b2b9c1624
commit
1397e01735
3 changed files with 26 additions and 14 deletions
|
@ -43,6 +43,20 @@ pub enum Relation {
|
||||||
Contributor,
|
Contributor,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl Model {
|
||||||
|
/// Returns the timestamp of when the user's account was created.
|
||||||
|
///
|
||||||
|
/// This will be the earlier of the `created_at` and `github_user_created_at` timestamps.
|
||||||
|
pub fn account_created_at(&self) -> NaiveDateTime {
|
||||||
|
let mut account_created_at = self.created_at;
|
||||||
|
if let Some(github_created_at) = self.github_user_created_at {
|
||||||
|
account_created_at = account_created_at.min(github_created_at);
|
||||||
|
}
|
||||||
|
|
||||||
|
account_created_at
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
impl Related<super::access_token::Entity> for Entity {
|
impl Related<super::access_token::Entity> for Entity {
|
||||||
fn to() -> RelationDef {
|
fn to() -> RelationDef {
|
||||||
Relation::AccessToken.def()
|
Relation::AccessToken.def()
|
||||||
|
|
|
@ -37,8 +37,7 @@ impl LlmTokenClaims {
|
||||||
user: &user::Model,
|
user: &user::Model,
|
||||||
is_staff: bool,
|
is_staff: bool,
|
||||||
billing_preferences: Option<billing_preference::Model>,
|
billing_preferences: Option<billing_preference::Model>,
|
||||||
has_llm_closed_beta_feature_flag: bool,
|
feature_flags: &Vec<String>,
|
||||||
has_predict_edits_feature_flag: bool,
|
|
||||||
has_llm_subscription: bool,
|
has_llm_subscription: bool,
|
||||||
plan: rpc::proto::Plan,
|
plan: rpc::proto::Plan,
|
||||||
system_id: Option<String>,
|
system_id: Option<String>,
|
||||||
|
@ -59,8 +58,12 @@ impl LlmTokenClaims {
|
||||||
metrics_id: user.metrics_id,
|
metrics_id: user.metrics_id,
|
||||||
github_user_login: user.github_login.clone(),
|
github_user_login: user.github_login.clone(),
|
||||||
is_staff,
|
is_staff,
|
||||||
has_llm_closed_beta_feature_flag,
|
has_llm_closed_beta_feature_flag: feature_flags
|
||||||
has_predict_edits_feature_flag,
|
.iter()
|
||||||
|
.any(|flag| flag == "llm-closed-beta"),
|
||||||
|
has_predict_edits_feature_flag: feature_flags
|
||||||
|
.iter()
|
||||||
|
.any(|flag| flag == "predict-edits"),
|
||||||
has_llm_subscription,
|
has_llm_subscription,
|
||||||
max_monthly_spend_in_cents: billing_preferences
|
max_monthly_spend_in_cents: billing_preferences
|
||||||
.map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {
|
.map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {
|
||||||
|
|
|
@ -4047,8 +4047,8 @@ async fn get_llm_api_token(
|
||||||
|
|
||||||
let flags = db.get_user_flags(session.user_id()).await?;
|
let flags = db.get_user_flags(session.user_id()).await?;
|
||||||
let has_language_models_feature_flag = flags.iter().any(|flag| flag == "language-models");
|
let has_language_models_feature_flag = flags.iter().any(|flag| flag == "language-models");
|
||||||
let has_llm_closed_beta_feature_flag = flags.iter().any(|flag| flag == "llm-closed-beta");
|
let has_bypass_account_age_check_feature_flag =
|
||||||
let has_predict_edits_feature_flag = flags.iter().any(|flag| flag == "predict-edits");
|
flags.iter().any(|flag| flag == "bypass-account-age-check");
|
||||||
|
|
||||||
if !session.is_staff() && !has_language_models_feature_flag {
|
if !session.is_staff() && !has_language_models_feature_flag {
|
||||||
Err(anyhow!("permission denied"))?
|
Err(anyhow!("permission denied"))?
|
||||||
|
@ -4067,13 +4067,9 @@ async fn get_llm_api_token(
|
||||||
let has_llm_subscription = session.has_llm_subscription(&db).await?;
|
let has_llm_subscription = session.has_llm_subscription(&db).await?;
|
||||||
|
|
||||||
let bypass_account_age_check =
|
let bypass_account_age_check =
|
||||||
has_llm_subscription || flags.iter().any(|flag| flag == "bypass-account-age-check");
|
has_llm_subscription || has_bypass_account_age_check_feature_flag;
|
||||||
if !bypass_account_age_check {
|
if !bypass_account_age_check {
|
||||||
let mut account_created_at = user.created_at;
|
if Utc::now().naive_utc() - user.account_created_at() < MIN_ACCOUNT_AGE_FOR_LLM_USE {
|
||||||
if let Some(github_created_at) = user.github_user_created_at {
|
|
||||||
account_created_at = account_created_at.min(github_created_at);
|
|
||||||
}
|
|
||||||
if Utc::now().naive_utc() - account_created_at < MIN_ACCOUNT_AGE_FOR_LLM_USE {
|
|
||||||
Err(anyhow!("account too young"))?
|
Err(anyhow!("account too young"))?
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -4084,8 +4080,7 @@ async fn get_llm_api_token(
|
||||||
&user,
|
&user,
|
||||||
session.is_staff(),
|
session.is_staff(),
|
||||||
billing_preferences,
|
billing_preferences,
|
||||||
has_llm_closed_beta_feature_flag,
|
&flags,
|
||||||
has_predict_edits_feature_flag,
|
|
||||||
has_llm_subscription,
|
has_llm_subscription,
|
||||||
session.current_plan(&db).await?,
|
session.current_plan(&db).await?,
|
||||||
session.system_id.clone(),
|
session.system_id.clone(),
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue