| #![allow(clippy::too_many_arguments)] |
|
|
| |
| |
| pub mod base62; |
| mod base64; |
| pub mod duration_fmt; |
| pub mod hex; |
| pub mod string_builder; |
|
|
| use ::base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; |
| use ::core::str::FromStr as _; |
| use ::prost::Message as _; |
| use ::reqwest::Client; |
| use ::std::borrow::Cow; |
| pub use base64::{from_base64, to_base64}; |
| pub use hex::{byte_to_hex, hex_to_byte}; |
| pub use string_builder::StringBuilder; |
|
|
| use super::model::userinfo::{ |
| GetTeamsResponse, ListActiveSessionsResponse, Session, StripeProfile, Team, UserProfile, |
| }; |
| use crate::{ |
| app::{ |
| lazy::{ |
| aggregated_usage_events_url, chat_models_url, filtered_usage_events_url, |
| is_on_new_pricing_url, server_config_url, teams_url, |
| }, |
| model::{ChainUsage, Checksum, DateTime, ExtToken, GcppHost, Hash, RawToken, Token}, |
| }, |
| common::model::userinfo::{MembershipType, SubscriptionStatus}, |
| core::{ |
| aiserver::v1::{ |
| AvailableModelsRequest, AvailableModelsResponse, GetAggregatedUsageEventsRequest, |
| GetAggregatedUsageEventsResponse, GetFilteredUsageEventsRequest, |
| GetFilteredUsageEventsResponse, GetServerConfigResponse, |
| }, |
| config::key_config, |
| }, |
| }; |
|
|
| mod sealed { |
| pub trait Sealed: Sized {} |
|
|
| impl Sealed for bool {} |
| impl Sealed for &'static str {} |
| impl Sealed for usize {} |
| } |
|
|
| pub trait ParseFromEnv: sealed::Sealed { |
| type Result = Self; |
| fn parse_from_env(key: &str, default: Self) -> Self::Result; |
| } |
|
|
| impl ParseFromEnv for bool { |
| #[inline] |
| fn parse_from_env(key: &str, default: Self) -> Self::Result { |
| ::std::env::var(key) |
| .ok() |
| .map(|mut val| { |
| let res = { |
| val.make_ascii_lowercase(); |
| val.trim() |
| }; |
| match res { |
| "true" | "1" => true, |
| "false" | "0" => false, |
| _ => default, |
| } |
| }) |
| .unwrap_or(default) |
| } |
| } |
|
|
| impl ParseFromEnv for &'static str { |
| type Result = Cow<'static, str>; |
| #[inline] |
| fn parse_from_env(key: &str, default: Self) -> Self::Result { |
| match ::std::env::var(key) { |
| Ok(mut value) => { |
| let trimmed = value.trim(); |
|
|
| if trimmed.is_empty() { |
| |
| Cow::Borrowed(default) |
| } else if trimmed.len() == value.len() { |
| |
| Cow::Owned(value) |
| } else { |
| |
| let trimmed_len = trimmed.len(); |
| let start_offset = trimmed.as_ptr() as usize - value.as_ptr() as usize; |
|
|
| unsafe { |
| let vec = value.as_mut_vec(); |
| if start_offset > 0 { |
| vec.copy_within(start_offset..start_offset + trimmed_len, 0); |
| } |
| vec.set_len(trimmed_len); |
| } |
|
|
| Cow::Owned(value) |
| } |
| } |
| Err(_) => Cow::Borrowed(default), |
| } |
| } |
| } |
|
|
| impl ParseFromEnv for usize { |
| #[inline] |
| fn parse_from_env(key: &str, default: Self) -> Self::Result { |
| ::std::env::var(key) |
| .ok() |
| .and_then(|v| v.trim().parse().ok()) |
| .unwrap_or(default) |
| } |
| } |
|
|
| #[inline] |
| pub fn parse_from_env<T: ParseFromEnv>(key: &str, default: T) -> T::Result { |
| ParseFromEnv::parse_from_env(key, default) |
| } |
|
|
| pub fn now_secs() -> u64 { |
| std::time::SystemTime::now() |
| .duration_since(std::time::UNIX_EPOCH) |
| .expect("system time before Unix epoch") |
| .as_secs() |
| } |
|
|
| const LEN: usize = 2; |
|
|
| pub trait TrimNewlines: Sized { |
| fn trim_leading_newlines(self) -> Self; |
| } |
|
|
| impl TrimNewlines for &str { |
| #[inline(always)] |
| fn trim_leading_newlines(self) -> Self { |
| let bytes = self.as_bytes(); |
| if bytes.len() >= LEN && bytes[0] == b'\n' && bytes[1] == b'\n' { |
| return unsafe { self.get_unchecked(LEN..) }; |
| } |
| self |
| } |
| } |
|
|
| impl TrimNewlines for String { |
| #[inline(always)] |
| fn trim_leading_newlines(mut self) -> Self { |
| let bytes = self.as_bytes(); |
| if bytes.len() >= LEN && bytes[0] == b'\n' && bytes[1] == b'\n' { |
| unsafe { |
| let vec = self.as_mut_vec(); |
| vec.drain(..LEN); |
| } |
| } |
| self |
| } |
| } |
|
|
| |
| |
| #[inline(never)] |
| pub async fn get_token_profile( |
| client: Client, |
| token: &Token, |
| maybe_token: Option<&Token>, |
| is_pri: bool, |
| include_user: bool, |
| include_sessions: bool, |
| ) -> ( |
| Option<UserProfile>, |
| Option<StripeProfile>, |
| Option<Vec<Session>>, |
| ) { |
| let maybe_token = maybe_token.unwrap_or(token); |
|
|
| let mut buf = [0; 31]; |
| let user_id = maybe_token.raw().subject.id.to_str(&mut buf) as &str; |
|
|
| if include_user { |
| if include_sessions { |
| |
| let (mut stripe, mut user, teams, is_on_new_pricing, sessions) = tokio::join!( |
| get_stripe_profile(&client, token.as_str(), is_pri), |
| get_user_profile(&client, user_id, maybe_token.as_str(), is_pri), |
| get_teams(&client, user_id, maybe_token.as_str(), is_pri), |
| get_is_on_new_pricing(&client, user_id, maybe_token.as_str(), is_pri), |
| get_sessions(&client, user_id, maybe_token.as_str(), is_pri) |
| ); |
|
|
| if let Some(stripe) = stripe.as_mut() |
| && teams.is_some_and(|teams| { |
| teams.into_iter().any(|team| { |
| team.has_billing |
| && team.subscription_status.is_some_and(|subscription_status| { |
| matches!(subscription_status, SubscriptionStatus::Active) |
| }) |
| }) |
| }) |
| { |
| stripe.membership_type = MembershipType::Enterprise; |
| } |
|
|
| if let Some(user) = user.as_mut() { |
| user.is_on_new_pricing = is_on_new_pricing.unwrap_or(true); |
| } |
|
|
| |
| (user, stripe, sessions) |
| } else { |
| |
| let (mut stripe, mut user, teams, is_on_new_pricing) = tokio::join!( |
| get_stripe_profile(&client, token.as_str(), is_pri), |
| get_user_profile(&client, user_id, maybe_token.as_str(), is_pri), |
| get_teams(&client, user_id, maybe_token.as_str(), is_pri), |
| get_is_on_new_pricing(&client, user_id, maybe_token.as_str(), is_pri) |
| ); |
|
|
| if let Some(stripe) = stripe.as_mut() |
| && teams.is_some_and(|teams| { |
| teams.into_iter().any(|team| { |
| team.has_billing |
| && team.subscription_status.is_some_and(|subscription_status| { |
| matches!(subscription_status, SubscriptionStatus::Active) |
| }) |
| }) |
| }) |
| { |
| stripe.membership_type = MembershipType::Enterprise; |
| } |
|
|
| if let Some(user) = user.as_mut() { |
| user.is_on_new_pricing = is_on_new_pricing.unwrap_or(true); |
| } |
|
|
| |
| (user, stripe, None) |
| } |
| } else { |
| |
| let (mut stripe, teams) = tokio::join!( |
| get_stripe_profile(&client, token.as_str(), is_pri), |
| get_teams(&client, user_id, maybe_token.as_str(), is_pri), |
| ); |
|
|
| if let Some(stripe) = stripe.as_mut() |
| && teams.is_some_and(|teams| { |
| teams.into_iter().any(|team| { |
| team.has_billing |
| && team.subscription_status.is_some_and(|subscription_status| { |
| matches!(subscription_status, SubscriptionStatus::Active) |
| }) |
| }) |
| }) |
| { |
| stripe.membership_type = MembershipType::Enterprise; |
| } |
|
|
| (None, stripe, None) |
| } |
| } |
|
|
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| pub async fn get_stripe_profile( |
| client: &Client, |
| auth_token: &str, |
| is_pri: bool, |
| ) -> Option<StripeProfile> { |
| let request = super::client::build_profile_request(client, auth_token, is_pri); |
|
|
| let response = request.send().await.ok()?; |
| crate::debug!("<get_stripe_profile> {}", response.status()); |
| response.json::<StripeProfile>().await.ok() |
| } |
|
|
| |
| pub async fn get_user_profile( |
| client: &Client, |
| user_id: &str, |
| auth_token: &str, |
| is_pri: bool, |
| ) -> Option<UserProfile> { |
| let request = super::client::build_userinfo_request(client, user_id, auth_token, is_pri); |
|
|
| |
| |
| |
| |
| |
| let response = request.send().await.ok()?; |
| crate::debug!("<get_user_profile> {}", response.status()); |
| response.json::<UserProfile>().await.ok() |
| } |
|
|
| pub async fn get_available_models( |
| ext_token: ExtToken, |
| is_pri: bool, |
| mut request: AvailableModelsRequest, |
| ) -> Option<AvailableModelsResponse> { |
| request.exclude_max_named_models = true; |
| let response = { |
| let client = super::client::build_client_request(super::client::AiServiceRequest { |
| ext_token, |
| fs_client_key: None, |
| url: chat_models_url(is_pri), |
| is_stream: false, |
| trace_id: Some(new_uuid_v4()), |
| is_pri, |
| cookie: None, |
| }); |
| client |
| .body(__unwrap!(encode_message(&request, false))) |
| .send() |
| .await |
| .ok()? |
| .bytes() |
| .await |
| .ok()? |
| }; |
| let available_models = AvailableModelsResponse::decode(response.as_ref()).ok()?; |
| Some(available_models) |
| } |
|
|
| pub async fn get_token_usage( |
| ext_token: ExtToken, |
| is_pri: bool, |
| time: DateTime, |
| model_id: &'static str, |
| ) -> Option<ChainUsage> { |
| const POLL_MAX_ATTEMPTS: u8 = 4; |
| const POLL_INTERVAL: ::core::time::Duration = ::core::time::Duration::from_millis(1500); |
|
|
| let maybe_token = ext_token |
| .secondary_token |
| .as_ref() |
| .unwrap_or(&ext_token.primary_token); |
|
|
| let mut buf = [0; 31]; |
| let user_id = maybe_token.raw().subject.id.to_str(&mut buf) as &str; |
| let mut token_usage = None; |
|
|
| for _ in 0..POLL_MAX_ATTEMPTS { |
| tokio::time::sleep(POLL_INTERVAL).await; |
| let res = get_filtered_usage_events( |
| &ext_token.get_client(), |
| user_id, |
| maybe_token.as_str(), |
| is_pri, |
| FilteredUsageArgs { |
| start: Some(time), |
| end: None, |
| model_id: Some(model_id), |
| size: Some(10), |
| }, |
| ) |
| .await?; |
|
|
| if let Some(first) = res.usage_events_display.first() |
| && let Some(usage) = first.token_usage |
| { |
| token_usage = Some(usage); |
| break; |
| }; |
| } |
|
|
| unsafe { ::core::mem::transmute(token_usage) } |
| } |
|
|
| |
| |
| |
|
|
| |
| |
| |
| |
|
|
| |
| |
|
|
| |
| |
| |
|
|
| |
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| |
| |
|
|
| |
| |
| |
| |
|
|
| |
| |
| |
| |
| |
|
|
| |
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
|
|
| #[inline(always)] |
| pub fn format_time_ms(seconds: f64) -> f64 { (seconds * 1000.0).round() / 1000.0 } |
|
|
| |
| #[inline] |
| pub fn token_to_tokeninfo( |
| token: RawToken, |
| checksum: Checksum, |
| client_key: Hash, |
| config_version: Option<uuid::Uuid>, |
| session_id: uuid::Uuid, |
| proxy_name: Option<String>, |
| timezone: Option<String>, |
| gcpp_host: Option<i32>, |
| ) -> key_config::TokenInfo { |
| key_config::TokenInfo { |
| token: Some(key_config::token_info::Token::from_raw(token)), |
| checksum: checksum.into_bytes().to_vec(), |
| client_key: client_key.into_bytes().to_vec(), |
| config_version: config_version.map(|v| v.into_bytes().to_vec()), |
| session_id: session_id.into_bytes().to_vec(), |
| proxy_name, |
| timezone, |
| gcpp_host, |
| } |
| } |
|
|
| |
| #[inline] |
| pub fn tokeninfo_to_token(info: key_config::TokenInfo) -> Option<ExtToken> { |
| let checksum = Checksum::from_bytes(info.checksum.try_into().ok()?); |
| let client_key = Hash::from_bytes(info.client_key.try_into().ok()?); |
| let config_version = info |
| .config_version |
| .and_then(|v| uuid::Uuid::from_slice(&v).ok()); |
| let session_id = uuid::Uuid::from_slice(&info.session_id).ok()?; |
| let timezone = info.timezone.and_then(|s| chrono_tz::Tz::from_str(&s).ok()); |
| let gcpp_host = info.gcpp_host.and_then(GcppHost::from_i32); |
| Some(ExtToken { |
| primary_token: Token::new(info.token?.into_raw()?, None), |
| secondary_token: None, |
| checksum, |
| client_key, |
| config_version, |
| session_id, |
| proxy: info.proxy_name, |
| timezone, |
| gcpp_host, |
| user: None, |
| }) |
| } |
|
|
| |
| #[inline] |
| fn compress_gzip(data: &[u8]) -> Result<Vec<u8>, ::std::io::Error> { |
| use ::std::io::Write as _; |
| use flate2::{Compression, write::GzEncoder}; |
|
|
| const LEVEL: Compression = Compression::new(6); |
|
|
| |
| let estimated_size = data.len() / 2 + 18; |
| let mut encoder = GzEncoder::new(Vec::with_capacity(estimated_size), LEVEL); |
| encoder.write_all(data)?; |
| encoder.finish() |
| } |
|
|
| #[inline(always)] |
| pub fn encode_message( |
| message: &impl ::prost::Message, |
| maybe_stream: bool, |
| ) -> Result<Vec<u8>, Box<dyn ::std::error::Error + Send + Sync>> { |
| const COMPRESSION_THRESHOLD: usize = 1024; |
| const LENGTH_OVERFLOW_MSG: &str = "Message length exceeds ~4 GiB"; |
|
|
| let estimated_size = message.encoded_len(); |
|
|
| if !maybe_stream { |
| let mut encoded = Vec::with_capacity(estimated_size); |
| message.encode_raw(&mut encoded); |
| return Ok(encoded); |
| } |
|
|
| |
| if estimated_size > u32::MAX as usize { |
| __cold_path!(); |
| return Err(LENGTH_OVERFLOW_MSG.into()); |
| } |
|
|
| use ::core::mem::MaybeUninit; |
|
|
| |
| let mut buffer = Vec::<MaybeUninit<u8>>::with_capacity(5 + estimated_size); |
|
|
| unsafe { |
| |
| buffer.set_len(5 + estimated_size); |
|
|
| |
| let header_ptr: *mut u8 = buffer.as_mut_ptr().cast(); |
| let body_ptr = header_ptr.add(5); |
| message.encode_raw(&mut ::core::slice::from_raw_parts_mut( |
| body_ptr, |
| estimated_size, |
| )); |
|
|
| |
| let (compression_flag, final_len) = if estimated_size >= COMPRESSION_THRESHOLD { |
| let body_slice = ::core::slice::from_raw_parts(body_ptr, estimated_size); |
| let compressed = compress_gzip(body_slice)?; |
| let compressed_len = compressed.len(); |
|
|
| if compressed_len < estimated_size { |
| |
| ::core::ptr::copy_nonoverlapping(compressed.as_ptr(), body_ptr, compressed_len); |
| |
| buffer.set_len(5 + compressed_len); |
| (0x01, compressed_len) |
| } else { |
| (0x00, estimated_size) |
| } |
| } else { |
| (0x00, estimated_size) |
| }; |
|
|
| |
| *header_ptr = compression_flag; |
|
|
| |
| ::core::hint::assert_unchecked(final_len <= u32::MAX as usize); |
| let len_bytes = (final_len as u32).to_be_bytes(); |
|
|
| |
| ::core::ptr::copy_nonoverlapping(len_bytes.as_ptr(), header_ptr.add(1), 4); |
|
|
| |
| Ok(::core::mem::transmute(buffer)) |
| } |
| } |
|
|
| |
| #[inline] |
| fn generate_pkce_pair() -> ([u8; 43], [u8; 43]) { |
| use ::core::mem::MaybeUninit; |
| use rand::TryRngCore as _; |
| use sha2::Digest as _; |
|
|
| |
| let mut verifier_bytes = MaybeUninit::<[u8; 32]>::uninit(); |
|
|
| |
| unsafe { |
| let bytes_ptr = verifier_bytes.as_mut_ptr().cast(); |
| let bytes_slice = ::core::slice::from_raw_parts_mut(bytes_ptr, 32); |
|
|
| rand::rngs::OsRng |
| .try_fill_bytes(bytes_slice) |
| .expect("获取系统安全随机数失败,这是一个严重错误!"); |
|
|
| |
| let verifier_bytes = verifier_bytes.assume_init(); |
|
|
| |
| let mut code_verifier = MaybeUninit::<[u8; 43]>::uninit(); |
| let verifier_ptr = code_verifier.as_mut_ptr().cast(); |
| let verifier_slice = ::core::slice::from_raw_parts_mut(verifier_ptr, 43); |
| __unwrap!(URL_SAFE_NO_PAD.encode_slice(verifier_bytes, verifier_slice)); |
| let code_verifier = code_verifier.assume_init(); |
|
|
| |
| let hash_result = sha2::Sha256::digest(code_verifier); |
| let mut code_challenge = MaybeUninit::<[u8; 43]>::uninit(); |
| let challenge_ptr = code_challenge.as_mut_ptr().cast(); |
| let challenge_slice = ::core::slice::from_raw_parts_mut(challenge_ptr, 43); |
| __unwrap!(URL_SAFE_NO_PAD.encode_slice(hash_result, challenge_slice)); |
| let code_challenge = code_challenge.assume_init(); |
|
|
| (code_verifier, code_challenge) |
| } |
| } |
|
|
| pub async fn get_new_token(ext_token: &mut ExtToken, is_pri: bool) -> bool { |
| let is_session = ext_token.primary_token.is_session(); |
|
|
| match if is_session { |
| refresh_token(ext_token, is_pri).await |
| } else { |
| upgrade_token(ext_token, is_pri).await |
| } { |
| Some((new_token, s)) => { |
| let tmp = Token::new(new_token, Some(s)); |
| if !is_session && ext_token.secondary_token.is_none() { |
| let old_token = ::core::mem::replace(&mut ext_token.primary_token, tmp); |
| ext_token.secondary_token = Some(old_token); |
| } else { |
| ext_token.primary_token = tmp; |
| } |
| true |
| } |
| None => false, |
| } |
| } |
|
|
| async fn upgrade_token(ext_token: &ExtToken, is_pri: bool) -> Option<(RawToken, String)> { |
| const POLL_MAX_ATTEMPTS: u8 = 5; |
| const POLL_INTERVAL: ::core::time::Duration = ::core::time::Duration::from_secs(1); |
|
|
| #[derive(::serde::Deserialize)] |
| #[serde(rename_all = "camelCase")] |
| struct PollResponse { |
| pub access_token: String, |
| |
| |
| |
| |
| } |
|
|
| let (verifier, challenge) = generate_pkce_pair(); |
| let verifier = unsafe { ::core::str::from_utf8_unchecked(&verifier) }; |
| let challenge = unsafe { ::core::str::from_utf8_unchecked(&challenge) }; |
| let mut buf = [0; 36]; |
| let uuid = uuid::Uuid::new_v4().hyphenated().encode_lower(&mut buf) as &str; |
|
|
| let token = ext_token |
| .secondary_token |
| .as_ref() |
| .unwrap_or(&ext_token.primary_token); |
| let mut buf = [0; 31]; |
| let user_id = token.raw().subject.id.to_str(&mut buf) as &str; |
| let auth_token = token.as_str(); |
|
|
| |
| let upgrade_response = super::client::build_token_upgrade_request( |
| &ext_token.get_client(), |
| uuid, |
| challenge, |
| user_id, |
| auth_token, |
| is_pri, |
| ) |
| .send() |
| .await |
| .ok()?; |
|
|
| if !upgrade_response.status().is_success() { |
| return None; |
| } |
|
|
| |
| for _ in 0..POLL_MAX_ATTEMPTS { |
| let poll_response = super::client::build_token_poll_request( |
| &ext_token.get_client(), |
| uuid, |
| verifier, |
| is_pri, |
| ) |
| .send() |
| .await |
| .ok()?; |
|
|
| match poll_response.status() { |
| reqwest::StatusCode::OK => { |
| let token = poll_response |
| .json::<PollResponse>() |
| .await |
| .ok()? |
| .access_token; |
| return parse_token(token); |
| } |
| reqwest::StatusCode::NOT_FOUND => { |
| tokio::time::sleep(POLL_INTERVAL).await; |
| } |
| _ => return None, |
| } |
| } |
|
|
| None |
| } |
|
|
| async fn refresh_token(ext_token: &ExtToken, is_pri: bool) -> Option<(RawToken, String)> { |
| const CLIENT_ID: &str = "KbZUR41cY7W6zRSdpSUJ7I7mLYBKOCmB"; |
|
|
| struct RefreshTokenRequest<'a> { |
| refresh_token: &'a str, |
| } |
|
|
| impl ::serde::Serialize for RefreshTokenRequest<'_> { |
| fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> |
| where |
| S: ::serde::Serializer, |
| { |
| use ::serde::ser::SerializeStruct as _; |
| let mut state = serializer.serialize_struct("RefreshTokenRequest", 3)?; |
| state.serialize_field("grant_type", "refresh_token")?; |
| state.serialize_field("client_id", CLIENT_ID)?; |
| state.serialize_field("refresh_token", self.refresh_token)?; |
| state.end() |
| } |
| } |
|
|
| #[derive(::serde::Deserialize)] |
| struct RefreshTokenResponse { |
| access_token: String, |
| |
| |
| |
| } |
|
|
| let refresh_request = RefreshTokenRequest { |
| refresh_token: ext_token.primary_token.as_str(), |
| }; |
|
|
| let body = serde_json::to_vec(&refresh_request).ok()?; |
|
|
| let response = |
| super::client::build_token_refresh_request(&ext_token.get_client(), is_pri, body) |
| .send() |
| .await |
| .ok()?; |
|
|
| let token = response |
| .json::<RefreshTokenResponse>() |
| .await |
| .ok()? |
| .access_token; |
|
|
| parse_token(token) |
| } |
|
|
| |
| #[inline] |
| fn parse_token(token_string: String) -> Option<(RawToken, String)> { |
| let raw_token = token_string.parse().ok()?; |
| Some((raw_token, token_string)) |
| } |
|
|
| pub async fn get_server_config(ext_token: ExtToken, is_pri: bool) -> Option<uuid::Uuid> { |
| let response = { |
| let client = super::client::build_client_request(super::client::AiServiceRequest { |
| ext_token, |
| fs_client_key: None, |
| url: server_config_url(is_pri), |
| is_stream: false, |
| trace_id: Some(new_uuid_v4()), |
| is_pri, |
| cookie: None, |
| }); |
| client.send().await.ok()?.bytes().await.ok()? |
| }; |
| let server_config = GetServerConfigResponse::decode(response.as_ref()).ok()?; |
| uuid::Uuid::try_parse(&server_config.config_version).ok() |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| const EMPTY_JSON: bytes::Bytes = bytes::Bytes::from_static(b"{}"); |
|
|
| pub async fn get_teams( |
| client: &Client, |
| user_id: &str, |
| auth_token: &str, |
| is_pri: bool, |
| ) -> Option<Vec<Team>> { |
| let request = super::client::build_proto_web_request( |
| client, user_id, auth_token, teams_url, is_pri, EMPTY_JSON, |
| ); |
|
|
| request |
| .send() |
| .await |
| .ok()? |
| .json::<GetTeamsResponse>() |
| .await |
| .ok() |
| .map(|r| r.teams) |
| } |
|
|
| pub async fn get_is_on_new_pricing( |
| client: &Client, |
| user_id: &str, |
| auth_token: &str, |
| is_pri: bool, |
| ) -> Option<bool> { |
| let request = super::client::build_proto_web_request( |
| client, |
| user_id, |
| auth_token, |
| is_on_new_pricing_url, |
| is_pri, |
| EMPTY_JSON, |
| ); |
|
|
| #[derive(serde::Deserialize)] |
| struct PricingConfig { |
| #[serde(rename = "isOnNewPricing")] |
| is_on_new_pricing: bool, |
| } |
|
|
| request |
| .send() |
| .await |
| .ok()? |
| .json::<PricingConfig>() |
| .await |
| .ok() |
| .map(|r| r.is_on_new_pricing) |
| } |
|
|
| pub async fn get_sessions( |
| client: &Client, |
| user_id: &str, |
| auth_token: &str, |
| is_pri: bool, |
| ) -> Option<Vec<Session>> { |
| let request = super::client::build_sessions_request(client, user_id, auth_token, is_pri); |
|
|
| request |
| .send() |
| .await |
| .ok()? |
| .json::<ListActiveSessionsResponse>() |
| .await |
| .ok() |
| .map(|r| r.sessions) |
| } |
|
|
| pub async fn get_aggregated_usage_events( |
| client: &Client, |
| user_id: &str, |
| auth_token: &str, |
| is_pri: bool, |
| ) -> Option<GetAggregatedUsageEventsResponse> { |
| let request = super::client::build_proto_web_request( |
| client, |
| user_id, |
| auth_token, |
| aggregated_usage_events_url, |
| is_pri, |
| bytes::Bytes::from(__unwrap!(serde_json::to_vec(&{ |
| const DELTA: chrono::TimeDelta = __unwrap!(chrono::TimeDelta::new(2629743, 765840000)); |
| let now = DateTime::utc_now(); |
| let start_date = now - DELTA; |
| GetAggregatedUsageEventsRequest { |
| team_id: -1, |
| start_date: Some(start_date.timestamp_millis()), |
| end_date: Some(now.timestamp_millis()), |
| user_id: None, |
| } |
| }))), |
| ); |
|
|
| request |
| .send() |
| .await |
| .ok()? |
| .json::<GetAggregatedUsageEventsResponse>() |
| .await |
| .ok() |
| } |
|
|
| pub struct FilteredUsageArgs { |
| pub start: Option<DateTime>, |
| pub end: Option<DateTime>, |
| pub model_id: Option<&'static str>, |
| pub size: Option<i32>, |
| } |
|
|
| impl From<FilteredUsageArgs> for GetFilteredUsageEventsRequest { |
| #[inline] |
| fn from(args: FilteredUsageArgs) -> Self { |
| const TZ: chrono::FixedOffset = __unwrap!(chrono::FixedOffset::west_opt(16 * 3600)); |
| const TIME: chrono::NaiveTime = __unwrap!(chrono::NaiveTime::from_hms_opt(0, 0, 0)); |
| const START: chrono::TimeDelta = chrono::TimeDelta::days(-7); |
| const END: chrono::TimeDelta = __unwrap!(chrono::TimeDelta::new(86399, 999000000)); |
|
|
| let (start_date, end_date) = if let (Some(a), Some(b)) = (args.start, args.end) { |
| (a.timestamp_millis(), b.timestamp_millis()) |
| } else { |
| let now = chrono::DateTime::<chrono::FixedOffset>::from_naive_utc_and_offset( |
| DateTime::naive_now(), |
| TZ, |
| ) |
| .date_naive() |
| .and_time(TIME); |
| match (args.start, args.end) { |
| (None, None) => ( |
| (now + START) |
| .and_local_timezone(TZ) |
| .unwrap() |
| .timestamp_millis(), |
| (now + END) |
| .and_local_timezone(TZ) |
| .unwrap() |
| .timestamp_millis(), |
| ), |
| (None, Some(b)) => ( |
| (now + START) |
| .and_local_timezone(TZ) |
| .unwrap() |
| .timestamp_millis(), |
| b.timestamp_millis(), |
| ), |
| (Some(a), None) => ( |
| a.timestamp_millis(), |
| (now + END) |
| .and_local_timezone(TZ) |
| .unwrap() |
| .timestamp_millis(), |
| ), |
| (Some(_), Some(_)) => unsafe { ::core::hint::unreachable_unchecked() }, |
| } |
| }; |
| Self { |
| team_id: 0, |
| start_date: Some(start_date), |
| end_date: Some(end_date), |
| user_id: None, |
| model_id: args.model_id.map(ToString::to_string), |
| page: Some(1), |
| page_size: Some(args.size.unwrap_or(100)), |
| } |
| } |
| } |
|
|
| pub async fn get_filtered_usage_events( |
| client: &Client, |
| user_id: &str, |
| auth_token: &str, |
| is_pri: bool, |
| args: FilteredUsageArgs, |
| ) -> Option<GetFilteredUsageEventsResponse> { |
| let request = super::client::build_proto_web_request( |
| client, |
| user_id, |
| auth_token, |
| filtered_usage_events_url, |
| is_pri, |
| bytes::Bytes::from(__unwrap!(serde_json::to_vec(&{ |
| let req: GetFilteredUsageEventsRequest = args.into(); |
| req |
| }))), |
| ); |
|
|
| let res = request.send().await.ok()?; |
| crate::debug!("<get_filtered_usage_events> {}", res.status()); |
| let res = res.bytes().await.ok()?; |
| crate::debug!("<get_filtered_usage_events> {}", unsafe { |
| ::core::str::from_utf8_unchecked(&res[..]) |
| }); |
| serde_json::from_slice(&res[..]).ok() |
| |
| |
| |
| } |
|
|
| #[inline] |
| pub fn new_uuid_v4() -> [u8; 36] { |
| let mut buf = [0; 36]; |
| uuid::Uuid::new_v4().hyphenated().encode_lower(&mut buf); |
| buf |
| } |
|
|