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},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::LazyLock;
|
||||
|
@ -98,10 +99,15 @@ impl KeyCheckerConfig {
|
|||
.merge(Serialized::defaults(Cli::parse()))
|
||||
.extract()?;
|
||||
|
||||
dbg!(&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
|
||||
|
@ -114,6 +120,22 @@ static DEFAULT_CONFIG: LazyLock<KeyCheckerConfig> = LazyLock::new(|| KeyCheckerC
|
|||
concurrency: 50,
|
||||
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 {
|
||||
DEFAULT_CONFIG.api_host.clone()
|
||||
}
|
||||
|
|
|
@ -1,5 +1,5 @@
|
|||
mod basic_config;
|
||||
mod basic_client;
|
||||
|
||||
pub use basic_config::KeyCheckerConfig;
|
||||
pub use basic_config::{KeyCheckerConfig, TEST_MESSAGE_BODY};
|
||||
pub use basic_client::client_builder;
|
||||
|
|
|
@ -1,15 +1,15 @@
|
|||
use anyhow::Result;
|
||||
use backon::{ExponentialBuilder, Retryable};
|
||||
use reqwest::{Client, StatusCode};
|
||||
use serde_json;
|
||||
use tokio::time::Duration;
|
||||
use url::Url;
|
||||
|
||||
use crate::config::TEST_MESSAGE_BODY;
|
||||
use crate::types::{GeminiKey, KeyStatus};
|
||||
|
||||
pub async fn validate_key_with_retry(
|
||||
client: Client,
|
||||
api_host: Url,
|
||||
full_url: Url,
|
||||
key: GeminiKey,
|
||||
) -> Option<GeminiKey> {
|
||||
let retry_policy = ExponentialBuilder::default()
|
||||
|
@ -17,7 +17,7 @@ pub async fn validate_key_with_retry(
|
|||
.with_min_delay(Duration::from_secs(3))
|
||||
.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) => {
|
||||
println!("Key: {}... -> SUCCESS", &key.as_ref()[..10]);
|
||||
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> {
|
||||
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"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
async fn keytest(client: Client, full_url: &Url, key: &GeminiKey) -> Result<KeyStatus> {
|
||||
let response = client
|
||||
.post(full_url)
|
||||
.post(full_url.clone())
|
||||
.header("Content-Type", "application/json")
|
||||
.header("X-goog-api-key", key.as_ref())
|
||||
.json(&request_body)
|
||||
.json(&*TEST_MESSAGE_BODY)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
|
|
|
@ -13,11 +13,17 @@ use crate::types::GeminiKey;
|
|||
pub struct ValidationService {
|
||||
config: KeyCheckerConfig,
|
||||
client: Client,
|
||||
full_url: url::Url,
|
||||
}
|
||||
|
||||
impl ValidationService {
|
||||
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<()> {
|
||||
|
@ -42,7 +48,7 @@ impl ValidationService {
|
|||
|
||||
// Create stream to validate keys concurrently
|
||||
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)
|
||||
.filter_map(|r| async { r });
|
||||
pin_mut!(valid_keys_stream);
|
||||
|
|
Loading…
Reference in New Issue