refactored every component thoroughly.
This commit is contained in:
glitchySid
262
src/gemini/client.rs
Normal file
262
src/gemini/client.rs
Normal file
@@ -0,0 +1,262 @@
|
||||
use crate::models::OrganizationPlan;
|
||||
use crate::storage::Cache;
|
||||
use crate::gemini::errors::GeminiError;
|
||||
use crate::gemini::prompt::PromptBuilder;
|
||||
use crate::gemini::types::{GeminiResponse, OrganizationPlanResponse};
|
||||
use reqwest::Client;
|
||||
use serde_json::json;
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
const DEFAULT_MODEL: &str = "gemini-3-flash-preview";
|
||||
const DEFAULT_TIMEOUT_SECS: u64 = 30;
|
||||
const MAX_RETRIES: u32 = 3;
|
||||
|
||||
pub struct GeminiClient {
|
||||
api_key: String,
|
||||
client: Client,
|
||||
base_url: String,
|
||||
#[allow(dead_code)]
|
||||
model: String,
|
||||
#[allow(dead_code)]
|
||||
timeout: Duration,
|
||||
}
|
||||
|
||||
impl GeminiClient {
|
||||
pub fn new(api_key: String) -> Self {
|
||||
Self::with_model(api_key, DEFAULT_MODEL.to_string())
|
||||
}
|
||||
|
||||
pub fn with_model(api_key: String, model: String) -> Self {
|
||||
let timeout = Duration::from_secs(DEFAULT_TIMEOUT_SECS);
|
||||
let client = Self::build_client(timeout);
|
||||
let base_url = Self::build_base_url(&model);
|
||||
|
||||
Self {
|
||||
api_key,
|
||||
client,
|
||||
base_url,
|
||||
model,
|
||||
timeout,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_client(timeout: Duration) -> Client {
|
||||
Client::builder()
|
||||
.timeout(timeout)
|
||||
.build()
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn build_base_url(model: &str) -> String {
|
||||
format!(
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent",
|
||||
model
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn organize_files(
|
||||
&self,
|
||||
filenames: Vec<String>,
|
||||
) -> Result<OrganizationPlan, GeminiError> {
|
||||
self.organize_files_with_cache(filenames, None, None).await
|
||||
}
|
||||
|
||||
pub async fn organize_files_with_cache(
|
||||
&self,
|
||||
filenames: Vec<String>,
|
||||
mut cache: Option<&mut Cache>,
|
||||
base_path: Option<&Path>,
|
||||
) -> Result<OrganizationPlan, GeminiError> {
|
||||
let url = self.build_url();
|
||||
|
||||
if let (Some(cache), Some(base_path)) = (cache.as_ref(), base_path)
|
||||
&& let Some(cached_response) = cache.get_cached_response(&filenames, base_path)
|
||||
{
|
||||
return Ok(cached_response);
|
||||
}
|
||||
|
||||
let prompt = PromptBuilder::new(filenames.clone()).build_categorization_prompt();
|
||||
let request_body = self.build_categorization_request(&prompt);
|
||||
|
||||
let res = self.send_request_with_retry(&url, &request_body).await?;
|
||||
let plan = self.parse_categorization_response(res).await?;
|
||||
|
||||
if let (Some(cache), Some(base_path)) = (cache.as_mut(), base_path) {
|
||||
cache.cache_response(&filenames, plan.clone(), base_path);
|
||||
}
|
||||
|
||||
Ok(plan)
|
||||
}
|
||||
|
||||
fn build_url(&self) -> String {
|
||||
format!("{}?key={}", self.base_url, self.api_key)
|
||||
}
|
||||
|
||||
fn build_categorization_request(&self, prompt: &str) -> serde_json::Value {
|
||||
json!({
|
||||
"contents": [{ "parts": [{ "text": prompt }] }],
|
||||
"generationConfig": { "response_mime_type": "application/json" }
|
||||
})
|
||||
}
|
||||
|
||||
async fn parse_categorization_response(
|
||||
&self,
|
||||
res: reqwest::Response,
|
||||
) -> Result<OrganizationPlan, GeminiError> {
|
||||
if !res.status().is_success() {
|
||||
return Err(GeminiError::from_response(res).await);
|
||||
}
|
||||
|
||||
let gemini_response: GeminiResponse =
|
||||
res.json().await.map_err(GeminiError::NetworkError)?;
|
||||
|
||||
let raw_text = self.extract_text_from_response(&gemini_response)?;
|
||||
let plan_response: OrganizationPlanResponse = serde_json::from_str(&raw_text)?;
|
||||
|
||||
Ok(plan_response.to_organization_plan())
|
||||
}
|
||||
|
||||
fn extract_text_from_response(&self, response: &GeminiResponse) -> Result<String, GeminiError> {
|
||||
response
|
||||
.candidates
|
||||
.first()
|
||||
.ok_or_else(|| GeminiError::InvalidResponse("No candidates in response".to_string()))?
|
||||
.content
|
||||
.parts
|
||||
.first()
|
||||
.ok_or_else(|| GeminiError::InvalidResponse("No parts in content".to_string()))
|
||||
.map(|p| p.text.clone())
|
||||
}
|
||||
|
||||
async fn send_request_with_retry(
|
||||
&self,
|
||||
url: &str,
|
||||
request_body: &serde_json::Value,
|
||||
) -> Result<reqwest::Response, GeminiError> {
|
||||
let mut attempts = 0;
|
||||
let mut base_delay = Duration::from_secs(2);
|
||||
|
||||
loop {
|
||||
attempts += 1;
|
||||
|
||||
match self.client.post(url).json(request_body).send().await {
|
||||
Ok(response) => {
|
||||
if response.status().is_success() {
|
||||
return Ok(response);
|
||||
}
|
||||
|
||||
let error = GeminiError::from_response(response).await;
|
||||
|
||||
if error.is_retryable() && attempts < MAX_RETRIES {
|
||||
let delay = error.retry_delay().unwrap_or(base_delay);
|
||||
self.print_retry_message(&error, delay, attempts);
|
||||
tokio::time::sleep(delay).await;
|
||||
base_delay *= 2;
|
||||
continue;
|
||||
}
|
||||
|
||||
return Err(error);
|
||||
}
|
||||
Err(e) => {
|
||||
if attempts < MAX_RETRIES {
|
||||
self.print_network_retry(&e, base_delay, attempts);
|
||||
tokio::time::sleep(base_delay).await;
|
||||
base_delay *= 2;
|
||||
continue;
|
||||
}
|
||||
return Err(GeminiError::NetworkError(e));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn print_retry_message(&self, error: &GeminiError, delay: Duration, attempt: u32) {
|
||||
println!(
|
||||
"API Error: {}. Retrying in {} seconds (attempt {}/{})",
|
||||
error,
|
||||
delay.as_secs(),
|
||||
attempt,
|
||||
MAX_RETRIES
|
||||
);
|
||||
}
|
||||
|
||||
fn print_network_retry(&self, error: &reqwest::Error, delay: Duration, attempt: u32) {
|
||||
println!(
|
||||
"Network error: {}. Retrying in {} seconds (attempt {}/{})",
|
||||
error,
|
||||
delay.as_secs(),
|
||||
attempt,
|
||||
MAX_RETRIES
|
||||
);
|
||||
}
|
||||
|
||||
pub async fn get_ai_sub_category(
|
||||
&self,
|
||||
filename: &str,
|
||||
parent_category: &str,
|
||||
content: &str,
|
||||
) -> String {
|
||||
let url = self.build_url();
|
||||
let prompt = PromptBuilder::build_subcategory_prompt(filename, parent_category, content);
|
||||
let request_body = self.build_subcategory_request(&prompt);
|
||||
|
||||
let res = match self.client.post(&url).json(&request_body).send().await {
|
||||
Ok(res) => res,
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"Warning: Failed to get sub-category for {}: {}",
|
||||
filename, e
|
||||
);
|
||||
return "General".to_string();
|
||||
}
|
||||
};
|
||||
|
||||
self.parse_subcategory_response(res, filename).await
|
||||
}
|
||||
|
||||
fn build_subcategory_request(&self, prompt: &str) -> serde_json::Value {
|
||||
json!({
|
||||
"contents": [{ "parts": [{ "text": prompt }] }]
|
||||
})
|
||||
}
|
||||
|
||||
async fn parse_subcategory_response(&self, res: reqwest::Response, filename: &str) -> String {
|
||||
if !res.status().is_success() {
|
||||
eprintln!(
|
||||
"Warning: API returned error for {}: {}",
|
||||
filename,
|
||||
res.status()
|
||||
);
|
||||
return "General".to_string();
|
||||
}
|
||||
|
||||
let gemini_response: GeminiResponse = match res.json().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
eprintln!("Warning: Failed to parse response for {}: {}", filename, e);
|
||||
return "General".to_string();
|
||||
}
|
||||
};
|
||||
|
||||
self.extract_subcategory_from_response(&gemini_response, filename)
|
||||
}
|
||||
|
||||
fn extract_subcategory_from_response(
|
||||
&self,
|
||||
response: &GeminiResponse,
|
||||
_filename: &str,
|
||||
) -> String {
|
||||
match self.extract_text_from_response(response) {
|
||||
Ok(text) => {
|
||||
let sub_category = text.trim();
|
||||
if sub_category.is_empty() {
|
||||
"General".to_string()
|
||||
} else {
|
||||
sub_category.to_string()
|
||||
}
|
||||
}
|
||||
Err(_) => "General".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user