feat: Add comprehensive improvements - CLI, error handling, and docs
- Add CLI argument parsing with clap (dry-run, max-concurrent options) - Replace .env configuration with interactive prompts and TOML config - Add BaseDirs-based configuration storage in ~/.config/noentropy/ - Improve Gemini API client with configurable model and timeout - Add concurrent processing with semaphore for rate limiting - Improve error handling with retry logic and exponential backoff - Add comprehensive README with installation and usage instructions - Add config.example.toml template for users - Update main.rs with better UX and colored output - Add lib.rs exports for config module - Refactor error response parsing for cleaner code - Update API endpoint to use configurable model parameter - Add proper error type handling in gemini_errors.rs
This commit is contained in:
129
src/gemini.rs
129
src/gemini.rs
@@ -42,14 +42,28 @@ pub struct GeminiClient {
|
||||
api_key: String,
|
||||
client: Client,
|
||||
base_url: String,
|
||||
model: String,
|
||||
timeout: Duration,
|
||||
}
|
||||
|
||||
impl GeminiClient {
|
||||
pub fn new(api_key: String) -> Self {
|
||||
Self::with_model(api_key, "gemini-3-flash-preview".to_string())
|
||||
}
|
||||
|
||||
pub fn with_model(api_key: String, model: String) -> Self {
|
||||
Self {
|
||||
api_key,
|
||||
client: Client::new(),
|
||||
base_url: "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent".to_string(),
|
||||
client: Client::builder()
|
||||
.timeout(Duration::from_secs(30))
|
||||
.build()
|
||||
.unwrap_or_default(),
|
||||
base_url: format!(
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent",
|
||||
model
|
||||
),
|
||||
model,
|
||||
timeout: Duration::from_secs(30),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -71,10 +85,10 @@ impl GeminiClient {
|
||||
let url = format!("{}?key={}", self.base_url, self.api_key);
|
||||
|
||||
// Check cache first if available
|
||||
if let (Some(cache_ref), Some(base_path)) = (cache.as_ref(), base_path) {
|
||||
if let Some(cached_response) = cache_ref.get_cached_response(&filenames, base_path) {
|
||||
return Ok(cached_response);
|
||||
}
|
||||
if let (Some(cache_ref), Some(base_path)) = (cache.as_ref(), base_path)
|
||||
&& let Some(cached_response) = cache_ref.get_cached_response(&filenames, base_path)
|
||||
{
|
||||
return Ok(cached_response);
|
||||
}
|
||||
|
||||
// 1. Construct the Prompt
|
||||
@@ -101,14 +115,19 @@ impl GeminiClient {
|
||||
|
||||
// 4. Parse
|
||||
if res.status().is_success() {
|
||||
let gemini_response: GeminiResponse = res.json().await.map_err(GeminiError::NetworkError)?;
|
||||
let gemini_response: GeminiResponse =
|
||||
res.json().await.map_err(GeminiError::NetworkError)?;
|
||||
|
||||
// Extract raw JSON string from Gemini using proper structs
|
||||
let raw_text = &gemini_response.candidates
|
||||
.get(0)
|
||||
.ok_or_else(|| GeminiError::InvalidResponse("No candidates in response".to_string()))?
|
||||
.content.parts
|
||||
.get(0)
|
||||
let raw_text = &gemini_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()))?
|
||||
.text;
|
||||
|
||||
@@ -147,6 +166,7 @@ impl GeminiClient {
|
||||
) -> Result<reqwest::Response, GeminiError> {
|
||||
let mut attempts = 0;
|
||||
let max_attempts = 3;
|
||||
let mut base_delay = Duration::from_secs(2);
|
||||
|
||||
loop {
|
||||
attempts += 1;
|
||||
@@ -158,23 +178,34 @@ impl GeminiClient {
|
||||
}
|
||||
|
||||
let error = GeminiError::from_response(response).await;
|
||||
|
||||
|
||||
if error.is_retryable() && attempts < max_attempts {
|
||||
if let Some(delay) = error.retry_delay() {
|
||||
println!("API Error: {}. Retrying in {} seconds (attempt {}/{})",
|
||||
error, delay.as_secs(), attempts, max_attempts);
|
||||
tokio::time::sleep(delay).await;
|
||||
continue;
|
||||
}
|
||||
let delay = error.retry_delay().unwrap_or(base_delay);
|
||||
println!(
|
||||
"API Error: {}. Retrying in {} seconds (attempt {}/{})",
|
||||
error,
|
||||
delay.as_secs(),
|
||||
attempts,
|
||||
max_attempts
|
||||
);
|
||||
tokio::time::sleep(delay).await;
|
||||
base_delay *= 2;
|
||||
continue;
|
||||
}
|
||||
|
||||
|
||||
return Err(error);
|
||||
}
|
||||
Err(e) => {
|
||||
if attempts < max_attempts {
|
||||
println!("Network error: {}. Retrying in {} seconds (attempt {}/{})",
|
||||
e, 5, attempts, max_attempts);
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
println!(
|
||||
"Network error: {}. Retrying in {} seconds (attempt {}/{})",
|
||||
e,
|
||||
base_delay.as_secs(),
|
||||
attempts,
|
||||
max_attempts
|
||||
);
|
||||
tokio::time::sleep(base_delay).await;
|
||||
base_delay *= 2;
|
||||
continue;
|
||||
}
|
||||
return Err(GeminiError::NetworkError(e));
|
||||
@@ -202,28 +233,46 @@ impl GeminiClient {
|
||||
}]
|
||||
});
|
||||
|
||||
let res = self.client.post(&url).json(&request_body).send().await;
|
||||
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();
|
||||
}
|
||||
};
|
||||
|
||||
if let Ok(res) = res {
|
||||
if res.status().is_success() {
|
||||
let gemini_response: GeminiResponse = res.json().await.unwrap_or_default();
|
||||
let sub_category = gemini_response.candidates
|
||||
.get(0)
|
||||
.and_then(|c| c.content.parts.get(0))
|
||||
.map(|p| p.text.trim())
|
||||
.unwrap_or("General")
|
||||
.to_string();
|
||||
|
||||
if sub_category.is_empty() {
|
||||
"General".to_string()
|
||||
} else {
|
||||
sub_category
|
||||
if res.status().is_success() {
|
||||
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();
|
||||
}
|
||||
} else {
|
||||
};
|
||||
|
||||
let sub_category = gemini_response
|
||||
.candidates
|
||||
.first()
|
||||
.and_then(|c| c.content.parts.first())
|
||||
.map(|p| p.text.trim())
|
||||
.unwrap_or("General")
|
||||
.to_string();
|
||||
|
||||
if sub_category.is_empty() {
|
||||
"General".to_string()
|
||||
} else {
|
||||
sub_category
|
||||
}
|
||||
} else {
|
||||
eprintln!(
|
||||
"Warning: API returned error for {}: {}",
|
||||
filename,
|
||||
res.status()
|
||||
);
|
||||
"General".to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -73,8 +73,7 @@ impl GeminiError {
|
||||
/// Parse HTTP response and convert to appropriate GeminiError
|
||||
pub async fn from_response(response: Response) -> Self {
|
||||
let status = response.status();
|
||||
|
||||
// Try to parse error response body
|
||||
|
||||
let error_text = match response.text().await {
|
||||
Ok(text) => text,
|
||||
Err(e) => {
|
||||
@@ -82,12 +81,10 @@ impl GeminiError {
|
||||
}
|
||||
};
|
||||
|
||||
// Try to parse structured error response
|
||||
if let Ok(gemini_error) = serde_json::from_str::<GeminiErrorResponse>(&error_text) {
|
||||
return Self::from_gemini_error(gemini_error.error, status.as_u16());
|
||||
}
|
||||
|
||||
// Fallback to HTTP status code based errors
|
||||
Self::from_status_code(status, &error_text)
|
||||
}
|
||||
|
||||
@@ -96,13 +93,11 @@ impl GeminiError {
|
||||
|
||||
match error_detail.status.as_str() {
|
||||
"RESOURCE_EXHAUSTED" => {
|
||||
if let Some(retry_info) = details.iter().find(|d| d.retry_delay.is_some()) {
|
||||
if let Some(retry_delay) = &retry_info.retry_delay {
|
||||
if let Ok(seconds) = retry_delay.parse::<u32>() {
|
||||
if let Some(retry_info) = details.iter().find(|d| d.retry_delay.is_some())
|
||||
&& let Some(retry_delay) = &retry_info.retry_delay
|
||||
&& let Ok(seconds) = retry_delay.parse::<u32>() {
|
||||
return GeminiError::RateLimitExceeded { retry_after: seconds };
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(quota_info) = details.iter().find(|d| d.quota_limit.is_some()) {
|
||||
let limit = quota_info.quota_limit.as_deref().unwrap_or("unknown");
|
||||
@@ -177,7 +172,7 @@ impl GeminiError {
|
||||
500 => GeminiError::InternalError {
|
||||
details: error_text.to_string()
|
||||
},
|
||||
502 | 503 | 504 => GeminiError::ServiceUnavailable {
|
||||
502..=504 => GeminiError::ServiceUnavailable {
|
||||
reason: error_text.to_string()
|
||||
},
|
||||
_ => GeminiError::ApiError {
|
||||
@@ -189,14 +184,14 @@ impl GeminiError {
|
||||
|
||||
/// Check if this error is retryable
|
||||
pub fn is_retryable(&self) -> bool {
|
||||
match self {
|
||||
GeminiError::RateLimitExceeded { .. } => true,
|
||||
GeminiError::ServiceUnavailable { .. } => true,
|
||||
GeminiError::Timeout { .. } => true,
|
||||
GeminiError::NetworkError(_) => true,
|
||||
GeminiError::InternalError { .. } => true,
|
||||
_ => false,
|
||||
}
|
||||
matches!(
|
||||
self,
|
||||
GeminiError::RateLimitExceeded { .. }
|
||||
| GeminiError::ServiceUnavailable { .. }
|
||||
| GeminiError::Timeout { .. }
|
||||
| GeminiError::NetworkError(_)
|
||||
| GeminiError::InternalError { .. }
|
||||
)
|
||||
}
|
||||
|
||||
/// Get retry delay for retryable errors
|
||||
@@ -217,10 +212,9 @@ impl GeminiError {
|
||||
fn extract_model_name(message: &str) -> String {
|
||||
// Try to extract model name from error message
|
||||
// Example: "Model 'gemini-1.5-flash' not found"
|
||||
if let Some(start) = message.find('\'') {
|
||||
if let Some(end) = message[start + 1..].find('\'') {
|
||||
if let Some(start) = message.find('\'')
|
||||
&& let Some(end) = message[start + 1..].find('\'') {
|
||||
return message[start + 1..start + 1 + end].to_string();
|
||||
}
|
||||
}
|
||||
"unknown".to_string()
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod cache;
|
||||
pub mod config;
|
||||
pub mod files;
|
||||
pub mod gemini;
|
||||
pub mod gemini_errors;
|
||||
|
||||
79
src/main.rs
79
src/main.rs
@@ -1,36 +1,46 @@
|
||||
use clap::Parser;
|
||||
use colored::*;
|
||||
use futures::future::join_all;
|
||||
use noentropy::cache::Cache;
|
||||
use noentropy::config;
|
||||
use noentropy::files::{FileBatch, OrganizationPlan, execute_move};
|
||||
use noentropy::gemini::GeminiClient;
|
||||
use noentropy::gemini_errors::GeminiError;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(author, version, about, long_about = None)]
|
||||
struct Args {
|
||||
#[arg(short, long, help = "Preview changes without moving files")]
|
||||
dry_run: bool,
|
||||
|
||||
#[arg(
|
||||
short,
|
||||
long,
|
||||
default_value_t = 5,
|
||||
help = "Maximum concurrent API requests"
|
||||
)]
|
||||
max_concurrent: usize,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
dotenv::dotenv().ok();
|
||||
let api_key = std::env::var("GEMINI_API_KEY")
|
||||
.map_err(|_| "GEMINI_API_KEY environment variable not set. Please set it in your .env file.")?;
|
||||
let download_path_var = std::env::var("DOWNLOAD_FOLDER")
|
||||
.map_err(|_| "DOWNLOAD_FOLDER environment variable not set. Please set it in your .env file.")?;
|
||||
let args = Args::parse();
|
||||
let api_key = config::get_or_prompt_api_key()?;
|
||||
let download_path = config::get_or_prompt_download_folder()?;
|
||||
|
||||
// 1. Setup
|
||||
let download_path: PathBuf = PathBuf::from(download_path_var.to_string());
|
||||
let client: GeminiClient = GeminiClient::new(api_key);
|
||||
|
||||
// Initialize cache
|
||||
|
||||
let cache_path = Path::new(".noentropy_cache.json");
|
||||
let mut cache = Cache::load_or_create(cache_path);
|
||||
|
||||
// Clean up old cache entries (older than 7 days)
|
||||
|
||||
cache.cleanup_old_entries(7 * 24 * 60 * 60);
|
||||
|
||||
// 2. Get Files
|
||||
let batch = FileBatch::from_path(download_path.clone());
|
||||
|
||||
if batch.filenames.is_empty() {
|
||||
println!("No files found to organize!");
|
||||
println!("{}", "No files found to organize!".yellow());
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
@@ -39,7 +49,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
batch.count()
|
||||
);
|
||||
|
||||
// 3. Call Gemini for Initial Categorization
|
||||
let mut plan: OrganizationPlan = match client
|
||||
.organize_files_with_cache(batch.filenames, Some(&mut cache), Some(&download_path))
|
||||
.await
|
||||
@@ -51,22 +60,26 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
}
|
||||
};
|
||||
|
||||
println!("Gemini Plan received! Performing deep inspection...");
|
||||
println!("{}", "Gemini Plan received! Performing deep inspection...".green());
|
||||
|
||||
// 4. Deep Inspection - Process files concurrently
|
||||
let client = Arc::new(client);
|
||||
|
||||
let tasks: Vec<_> = plan.files.iter_mut()
|
||||
let semaphore = Arc::new(tokio::sync::Semaphore::new(args.max_concurrent));
|
||||
|
||||
let tasks: Vec<_> = plan
|
||||
.files
|
||||
.iter_mut()
|
||||
.zip(batch.paths.iter())
|
||||
.map(|(file_category, path)| {
|
||||
let client = Arc::clone(&client);
|
||||
let filename = file_category.filename.clone();
|
||||
let category = file_category.category.clone();
|
||||
let path = path.clone();
|
||||
|
||||
let semaphore = Arc::clone(&semaphore);
|
||||
|
||||
async move {
|
||||
if noentropy::files::is_text_file(&path) {
|
||||
if let Some(content) = noentropy::files::read_file_sample(&path, 2000) {
|
||||
let _permit = semaphore.acquire().await.unwrap();
|
||||
if let Some(content) = noentropy::files::read_file_sample(&path, 5000) {
|
||||
println!("Reading content of {}...", filename.green());
|
||||
client.get_ai_sub_category(&filename, &category, &content).await
|
||||
} else {
|
||||
@@ -79,22 +92,26 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Wait for all concurrent tasks to complete
|
||||
let sub_categories = join_all(tasks).await;
|
||||
|
||||
// Apply the results back to the plan
|
||||
|
||||
for (file_category, sub_category) in plan.files.iter_mut().zip(sub_categories) {
|
||||
file_category.sub_category = sub_category;
|
||||
}
|
||||
|
||||
println!("Deep inspection complete! Moving Files.....");
|
||||
// 5. Execute
|
||||
execute_move(&download_path, plan);
|
||||
println!("Done!");
|
||||
|
||||
// Save cache before exiting
|
||||
println!("{}", "Deep inspection complete! Moving Files.....".green());
|
||||
|
||||
if args.dry_run {
|
||||
println!(
|
||||
"{} Dry run mode - skipping file moves.",
|
||||
"INFO:".cyan()
|
||||
);
|
||||
} else {
|
||||
execute_move(&download_path, plan);
|
||||
}
|
||||
println!("{}", "Done!".green().bold());
|
||||
|
||||
if let Err(e) = cache.save(cache_path) {
|
||||
println!("Warning: Failed to save cache: {}", e);
|
||||
eprintln!("Warning: Failed to save cache: {}", e);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
|
||||
Reference in New Issue
Block a user