feat: enhance validation service with full API URL and refactor key validation logic

main
Yoo1tic 2025-07-19 23:09:15 +08:00
parent 43e535ec15
commit db94d6d86c
4 changed files with 39 additions and 25 deletions

View File

@ -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()
}

View File

@ -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;

View File

@ -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?;

View File

@ -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);