feat: enhance validation service with full API URL and refactor key validation logic
parent
43e535ec15
commit
db94d6d86c
|
@ -5,6 +5,7 @@ use figment::{
|
||||||
providers::{Env, Format, Serialized, Toml},
|
providers::{Env, Format, Serialized, Toml},
|
||||||
};
|
};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::Value;
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
use std::sync::LazyLock;
|
use std::sync::LazyLock;
|
||||||
|
@ -98,10 +99,15 @@ impl KeyCheckerConfig {
|
||||||
.merge(Serialized::defaults(Cli::parse()))
|
.merge(Serialized::defaults(Cli::parse()))
|
||||||
.extract()?;
|
.extract()?;
|
||||||
|
|
||||||
dbg!(&config);
|
|
||||||
|
|
||||||
Ok(config)
|
Ok(config)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Returns the complete Gemini API URL for generateContent endpoint
|
||||||
|
pub fn gemini_api_url(&self) -> Url {
|
||||||
|
self.api_host
|
||||||
|
.join("v1beta/models/gemini-2.0-flash-exp:generateContent")
|
||||||
|
.expect("Failed to join API URL")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Single LazyLock for entire default configuration
|
// Single LazyLock for entire default configuration
|
||||||
|
@ -114,6 +120,22 @@ static DEFAULT_CONFIG: LazyLock<KeyCheckerConfig> = LazyLock::new(|| KeyCheckerC
|
||||||
concurrency: 50,
|
concurrency: 50,
|
||||||
proxy: None,
|
proxy: None,
|
||||||
});
|
});
|
||||||
|
|
||||||
|
// LazyLock for the test message body used in API key validation
|
||||||
|
pub static TEST_MESSAGE_BODY: LazyLock<Value> = LazyLock::new(|| {
|
||||||
|
serde_json::json!({
|
||||||
|
"contents": [
|
||||||
|
{
|
||||||
|
"parts": [
|
||||||
|
{
|
||||||
|
"text": "Hi"
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
]
|
||||||
|
})
|
||||||
|
});
|
||||||
|
|
||||||
fn default_api_host() -> Url {
|
fn default_api_host() -> Url {
|
||||||
DEFAULT_CONFIG.api_host.clone()
|
DEFAULT_CONFIG.api_host.clone()
|
||||||
}
|
}
|
||||||
|
|
|
@ -1,5 +1,5 @@
|
||||||
mod basic_config;
|
mod basic_config;
|
||||||
mod basic_client;
|
mod basic_client;
|
||||||
|
|
||||||
pub use basic_config::KeyCheckerConfig;
|
pub use basic_config::{KeyCheckerConfig, TEST_MESSAGE_BODY};
|
||||||
pub use basic_client::client_builder;
|
pub use basic_client::client_builder;
|
||||||
|
|
|
@ -1,15 +1,15 @@
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use backon::{ExponentialBuilder, Retryable};
|
use backon::{ExponentialBuilder, Retryable};
|
||||||
use reqwest::{Client, StatusCode};
|
use reqwest::{Client, StatusCode};
|
||||||
use serde_json;
|
|
||||||
use tokio::time::Duration;
|
use tokio::time::Duration;
|
||||||
use url::Url;
|
use url::Url;
|
||||||
|
|
||||||
|
use crate::config::TEST_MESSAGE_BODY;
|
||||||
use crate::types::{GeminiKey, KeyStatus};
|
use crate::types::{GeminiKey, KeyStatus};
|
||||||
|
|
||||||
pub async fn validate_key_with_retry(
|
pub async fn validate_key_with_retry(
|
||||||
client: Client,
|
client: Client,
|
||||||
api_host: Url,
|
full_url: Url,
|
||||||
key: GeminiKey,
|
key: GeminiKey,
|
||||||
) -> Option<GeminiKey> {
|
) -> Option<GeminiKey> {
|
||||||
let retry_policy = ExponentialBuilder::default()
|
let retry_policy = ExponentialBuilder::default()
|
||||||
|
@ -17,7 +17,7 @@ pub async fn validate_key_with_retry(
|
||||||
.with_min_delay(Duration::from_secs(3))
|
.with_min_delay(Duration::from_secs(3))
|
||||||
.with_max_delay(Duration::from_secs(5));
|
.with_max_delay(Duration::from_secs(5));
|
||||||
|
|
||||||
let result = (async || match keytest(client.to_owned(), &api_host, &key).await {
|
let result = (async || match keytest(client.to_owned(), &full_url, &key).await {
|
||||||
Ok(KeyStatus::Valid) => {
|
Ok(KeyStatus::Valid) => {
|
||||||
println!("Key: {}... -> SUCCESS", &key.as_ref()[..10]);
|
println!("Key: {}... -> SUCCESS", &key.as_ref()[..10]);
|
||||||
Ok(Some(key.clone()))
|
Ok(Some(key.clone()))
|
||||||
|
@ -58,26 +58,12 @@ pub async fn validate_key_with_retry(
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn keytest(client: Client, api_host: &Url, key: &GeminiKey) -> Result<KeyStatus> {
|
async fn keytest(client: Client, full_url: &Url, key: &GeminiKey) -> Result<KeyStatus> {
|
||||||
const API_PATH: &str = "v1beta/models/gemini-2.0-flash-exp:generateContent";
|
|
||||||
let full_url = api_host.join(API_PATH)?;
|
|
||||||
let request_body = serde_json::json!({
|
|
||||||
"contents": [
|
|
||||||
{
|
|
||||||
"parts": [
|
|
||||||
{
|
|
||||||
"text": "Hi"
|
|
||||||
}
|
|
||||||
]
|
|
||||||
}
|
|
||||||
]
|
|
||||||
});
|
|
||||||
|
|
||||||
let response = client
|
let response = client
|
||||||
.post(full_url)
|
.post(full_url.clone())
|
||||||
.header("Content-Type", "application/json")
|
.header("Content-Type", "application/json")
|
||||||
.header("X-goog-api-key", key.as_ref())
|
.header("X-goog-api-key", key.as_ref())
|
||||||
.json(&request_body)
|
.json(&*TEST_MESSAGE_BODY)
|
||||||
.send()
|
.send()
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
|
|
|
@ -13,11 +13,17 @@ use crate::types::GeminiKey;
|
||||||
pub struct ValidationService {
|
pub struct ValidationService {
|
||||||
config: KeyCheckerConfig,
|
config: KeyCheckerConfig,
|
||||||
client: Client,
|
client: Client,
|
||||||
|
full_url: url::Url,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ValidationService {
|
impl ValidationService {
|
||||||
pub fn new(config: KeyCheckerConfig, client: Client) -> Self {
|
pub fn new(config: KeyCheckerConfig, client: Client) -> Self {
|
||||||
Self { config, client }
|
let full_url = config.gemini_api_url();
|
||||||
|
Self {
|
||||||
|
config,
|
||||||
|
client,
|
||||||
|
full_url,
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn validate_keys(&self, keys: Vec<GeminiKey>) -> Result<()> {
|
pub async fn validate_keys(&self, keys: Vec<GeminiKey>) -> Result<()> {
|
||||||
|
@ -42,7 +48,7 @@ impl ValidationService {
|
||||||
|
|
||||||
// Create stream to validate keys concurrently
|
// Create stream to validate keys concurrently
|
||||||
let valid_keys_stream = stream
|
let valid_keys_stream = stream
|
||||||
.map(|key| validate_key_with_retry(self.client.to_owned(), self.config.api_host.clone(), key))
|
.map(|key| validate_key_with_retry(self.client.to_owned(), self.full_url.clone(), key))
|
||||||
.buffer_unordered(self.config.concurrency)
|
.buffer_unordered(self.config.concurrency)
|
||||||
.filter_map(|r| async { r });
|
.filter_map(|r| async { r });
|
||||||
pin_mut!(valid_keys_stream);
|
pin_mut!(valid_keys_stream);
|
||||||
|
|
Loading…
Reference in New Issue