refactor: Simplify codebase by extracting modules and helpers

Extract code into focused modules for better maintainability:

New modules:
- gemini_types.rs (32 lines) - Response type definitions
- gemini_helpers.rs (51 lines) - Prompt builder and conversion helpers
- prompt.rs (130 lines) - User input and validation logic

Refactored files:
- gemini.rs: 278 -> 259 lines (-19 lines)
  * Extract response parsing into helper methods
  * Extract request building into separate methods
  * Extract retry logic into dedicated functions
  * Use PromptBuilder for cleaner prompt construction

- config.rs: 275 -> 127 lines (-148 lines)
  * Extract all prompting logic to prompt.rs module
  * Simplify with Default trait for Config
  * Cleaner API methods

Benefits:
- Better separation of concerns
- Easier to test and maintain
- Clearer module boundaries
- Reduced nesting and complexity
- All 31 tests still passing
This commit is contained in:
2025-12-29 00:35:14 +05:30
parent 32c1f283ac
commit 1a72116b9d
8 changed files with 405 additions and 353 deletions

View File

@@ -1,11 +1,6 @@
use colored::*; use colored::*;
use directories::{BaseDirs, ProjectDirs};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::fs; use std::path::PathBuf;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
const MAX_RETRIES: u32 = 3;
#[derive(Debug, Serialize, Deserialize, Clone)] #[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Config { pub struct Config {
@@ -14,28 +9,14 @@ pub struct Config {
} }
impl Config { impl Config {
fn get_config_dir() -> Result<PathBuf, Box<dyn std::error::Error>> { pub fn load() -> Result<Self, Box<dyn std::error::Error>> {
if let Some(proj_dirs) = ProjectDirs::from("dev", "noentropy", "NoEntropy") {
let config_dir = proj_dirs.config_dir().to_path_buf();
fs::create_dir_all(&config_dir)?;
Ok(config_dir)
} else {
Err("Failed to determine config directory".into())
}
}
fn get_config_path() -> Result<PathBuf, Box<dyn std::error::Error>> {
Ok(Self::get_config_dir()?.join("config.toml"))
}
pub fn load() -> Result<Config, Box<dyn std::error::Error>> {
let config_path = Self::get_config_path()?; let config_path = Self::get_config_path()?;
if !config_path.exists() { if !config_path.exists() {
return Err("Config file not found".into()); return Err("Config file not found".into());
} }
let content = fs::read_to_string(&config_path)?; let content = std::fs::read_to_string(&config_path)?;
let config: Config = toml::from_str(&content)?; let config: Config = toml::from_str(&content)?;
if config.api_key.is_empty() { if config.api_key.is_empty() {
@@ -50,7 +31,7 @@ impl Config {
let toml_string = toml::to_string_pretty(self)?; let toml_string = toml::to_string_pretty(self)?;
fs::write(&config_path, toml_string)?; std::fs::write(&config_path, toml_string)?;
println!( println!(
"{} Configuration saved to {}", "{} Configuration saved to {}",
@@ -74,6 +55,21 @@ impl Config {
Err(_) => Err("Download folder not configured".into()), Err(_) => Err("Download folder not configured".into()),
} }
} }
fn get_config_dir() -> Result<PathBuf, Box<dyn std::error::Error>> {
match directories::ProjectDirs::from("dev", "noentropy", "NoEntropy") {
Some(proj_dirs) => {
let config_dir = proj_dirs.config_dir().to_path_buf();
std::fs::create_dir_all(&config_dir)?;
Ok(config_dir)
}
None => Err("Failed to determine config directory".into()),
}
}
fn get_config_path() -> Result<PathBuf, Box<dyn std::error::Error>> {
Ok(Self::get_config_dir()?.join("config.toml"))
}
} }
pub fn get_or_prompt_api_key() -> Result<String, Box<dyn std::error::Error>> { pub fn get_or_prompt_api_key() -> Result<String, Box<dyn std::error::Error>> {
@@ -87,17 +83,9 @@ pub fn get_or_prompt_api_key() -> Result<String, Box<dyn std::error::Error>> {
println!("{}", "🔑 NoEntropy Configuration".bold().cyan()); println!("{}", "🔑 NoEntropy Configuration".bold().cyan());
println!("{}", "─────────────────────────────".cyan()); println!("{}", "─────────────────────────────".cyan());
let api_key = prompt_api_key()?; let api_key = crate::prompt::Prompter::prompt_api_key()?;
let mut config = if let Ok(cfg) = Config::load() {
cfg
} else {
Config {
api_key: api_key.clone(),
download_folder: PathBuf::new(),
}
};
let mut config = Config::load().unwrap_or_default();
config.api_key = api_key.clone(); config.api_key = api_key.clone();
config.save()?; config.save()?;
@@ -115,17 +103,9 @@ pub fn get_or_prompt_download_folder() -> Result<PathBuf, Box<dyn std::error::Er
println!(); println!();
println!("{}", "📁 Download folder not configured.".yellow()); println!("{}", "📁 Download folder not configured.".yellow());
let folder_path = prompt_download_folder()?; let folder_path = crate::prompt::Prompter::prompt_download_folder()?;
let mut config = if let Ok(cfg) = Config::load() {
cfg
} else {
Config {
api_key: String::new(),
download_folder: folder_path.clone(),
}
};
let mut config = Config::load().unwrap_or_default();
config.download_folder = folder_path.clone(); config.download_folder = folder_path.clone();
config.save()?; config.save()?;
@@ -133,141 +113,13 @@ pub fn get_or_prompt_download_folder() -> Result<PathBuf, Box<dyn std::error::Er
Ok(folder_path) Ok(folder_path)
} }
fn prompt_api_key() -> Result<String, Box<dyn std::error::Error>> { impl Default for Config {
let mut attempts = 0; fn default() -> Self {
Self {
println!(); api_key: String::new(),
println!("Get your API key at: {}", "https://ai.google.dev/".cyan().underline()); download_folder: PathBuf::new(),
println!("Enter your API Key (starts with 'AIza'):");
while attempts < MAX_RETRIES {
print!("API Key: ");
io::stdout().flush()?;
let mut input = String::new();
io::stdin().read_line(&mut input)?;
let key = input.trim();
if validate_api_key(key) {
return Ok(key.to_string());
}
attempts += 1;
let remaining = MAX_RETRIES - attempts;
eprintln!(
"{} Invalid API key format. Must start with 'AIza' and be around 39 characters.",
"".red()
);
if remaining > 0 {
eprintln!("Try again ({} attempts remaining):", remaining);
} }
} }
Err("Max retries exceeded. Please run again with a valid API key.".into())
}
fn prompt_download_folder() -> Result<PathBuf, Box<dyn std::error::Error>> {
let default_path = get_default_downloads_folder();
let default_display = default_path.to_string_lossy();
let mut attempts = 0;
println!(
"Enter path to folder to organize (e.g., {}):",
default_display.yellow()
);
println!("Or press Enter to use default: {}", default_display.green());
println!("Folder path: ");
while attempts < MAX_RETRIES {
print!("> ");
io::stdout().flush()?;
let mut input = String::new();
io::stdin().read_line(&mut input)?;
let input = input.trim();
let path = if input.is_empty() {
default_path.clone()
} else {
let expanded = expand_home(input);
PathBuf::from(expanded)
};
if validate_folder_path(&path) {
return Ok(path);
}
attempts += 1;
let remaining = MAX_RETRIES - attempts;
eprintln!("{} Invalid folder path.", "".red());
if !path.exists() {
eprintln!(" Path does not exist: {}", path.display());
} else if !path.is_dir() {
eprintln!(" Path is not a directory: {}", path.display());
}
if remaining > 0 {
eprintln!("Try again ({} attempts remaining):", remaining);
println!("Folder path: ");
}
}
Err("Max retries exceeded. Please run again with a valid folder path.".into())
}
fn validate_api_key(key: &str) -> bool {
if key.is_empty() {
return false;
}
if !key.starts_with("AIza") {
return false;
}
if key.len() < 35 || key.len() > 50 {
return false;
}
true
}
fn validate_folder_path(path: &Path) -> bool {
if !path.exists() {
return false;
}
if !path.is_dir() {
return false;
}
true
}
fn get_default_downloads_folder() -> PathBuf {
if let Some(base_dirs) = BaseDirs::new() {
let home = base_dirs.home_dir();
return home.join("Downloads");
}
PathBuf::from("./Downloads")
}
fn expand_home(path: &str) -> String {
if path.starts_with("~/")
&& let Some(base_dirs) = BaseDirs::new()
{
let home = base_dirs.home_dir();
return path.replacen("~", &home.to_string_lossy(), 1);
}
path.to_string()
} }
#[cfg(test)] #[cfg(test)]

2
src/config_mod.rs.bak Normal file
View File

@@ -0,0 +1,2 @@
pub mod config;
pub mod prompt;

View File

@@ -1,4 +1,5 @@
use crate::config::*; use crate::config::*;
use std::path::Path;
#[test] #[test]
fn test_config_serialization() { fn test_config_serialization() {
@@ -17,37 +18,37 @@ fn test_config_serialization() {
#[test] #[test]
fn test_validate_api_key_valid() { fn test_validate_api_key_valid() {
assert!(validate_api_key("AIzaSyB1234567890123456789012345678")); assert!(crate::prompt::Prompter::validate_api_key("AIzaSyB1234567890123456789012345678"));
assert!(validate_api_key("AIzaSyB123456789012345678901234567890")); assert!(crate::prompt::Prompter::validate_api_key("AIzaSyB123456789012345678901234567890"));
} }
#[test] #[test]
fn test_validate_api_key_invalid() { fn test_validate_api_key_invalid() {
assert!(!validate_api_key("")); assert!(!crate::prompt::Prompter::validate_api_key(""));
assert!(!validate_api_key("invalid_key")); assert!(!crate::prompt::Prompter::validate_api_key("invalid_key"));
assert!(!validate_api_key("BizaSyB1234567890123456789012345678")); assert!(!crate::prompt::Prompter::validate_api_key("BizaSyB1234567890123456789012345678"));
assert!(!validate_api_key("short")); assert!(!crate::prompt::Prompter::validate_api_key("short"));
} }
#[test] #[test]
fn test_validate_folder_path_valid() { fn test_validate_folder_path_valid() {
let temp_dir = tempfile::tempdir().unwrap(); let temp_dir = tempfile::tempdir().unwrap();
assert!(validate_folder_path(temp_dir.path())); assert!(crate::prompt::Prompter::validate_folder_path(temp_dir.path()));
} }
#[test] #[test]
fn test_validate_folder_path_invalid() { fn test_validate_folder_path_invalid() {
assert!(!validate_folder_path(Path::new("/nonexistent/path/that/does/not/exist"))); assert!(!crate::prompt::Prompter::validate_folder_path(Path::new("/nonexistent/path/that/does/not/exist")));
let temp_file = tempfile::NamedTempFile::new().unwrap(); let temp_file = tempfile::NamedTempFile::new().unwrap();
assert!(!validate_folder_path(temp_file.path())); assert!(!crate::prompt::Prompter::validate_folder_path(temp_file.path()));
} }
#[test] #[test]
fn test_expand_home_with_tilde() { fn test_expand_home_with_tilde() {
if let Some(base_dirs) = BaseDirs::new() { if let Some(base_dirs) = directories::BaseDirs::new() {
let home = base_dirs.home_dir(); let home = base_dirs.home_dir();
let expanded = expand_home("~/test/path"); let expanded = crate::prompt::Prompter::expand_home("~/test/path");
assert!(expanded.starts_with(home.to_string_lossy().as_ref())); assert!(expanded.starts_with(home.to_string_lossy().as_ref()));
assert!(expanded.contains("test/path")); assert!(expanded.contains("test/path"));
} }
@@ -55,16 +56,16 @@ fn test_expand_home_with_tilde() {
#[test] #[test]
fn test_expand_home_without_tilde() { fn test_expand_home_without_tilde() {
let expanded = expand_home("/absolute/path"); let expanded = crate::prompt::Prompter::expand_home("/absolute/path");
assert_eq!(expanded, "/absolute/path"); assert_eq!(expanded, "/absolute/path");
let expanded = expand_home("relative/path"); let expanded = crate::prompt::Prompter::expand_home("relative/path");
assert_eq!(expanded, "relative/path"); assert_eq!(expanded, "relative/path");
} }
#[test] #[test]
fn test_get_default_downloads_folder() { fn test_get_default_downloads_folder() {
let path = get_default_downloads_folder(); let path = crate::prompt::Prompter::get_default_downloads_folder();
assert!(path.ends_with("Downloads")); assert!(path.ends_with("Downloads"));
} }

View File

@@ -1,42 +1,16 @@
use crate::cache::Cache; use crate::cache::Cache;
use crate::files::{FileCategory, OrganizationPlan};
use crate::gemini_errors::GeminiError; use crate::gemini_errors::GeminiError;
use crate::gemini_helpers::PromptBuilder;
use crate::gemini_types::{GeminiResponse, OrganizationPlanResponse};
use crate::files::OrganizationPlan;
use reqwest::Client; use reqwest::Client;
use serde::Deserialize;
use serde_json::json; use serde_json::json;
use std::path::Path; use std::path::Path;
use std::time::Duration; use std::time::Duration;
#[derive(Deserialize, Default)] const DEFAULT_MODEL: &str = "gemini-3-flash-preview";
struct GeminiResponse { const DEFAULT_TIMEOUT_SECS: u64 = 30;
candidates: Vec<Candidate>, const MAX_RETRIES: u32 = 3;
}
#[derive(Deserialize)]
struct Candidate {
content: Content,
}
#[derive(Deserialize)]
struct Content {
parts: Vec<Part>,
}
#[derive(Deserialize)]
struct Part {
text: String,
}
#[derive(Deserialize)]
struct FileCategoryResponse {
filename: String,
category: String,
}
#[derive(Deserialize)]
struct OrganizationPlanResponse {
files: Vec<FileCategoryResponse>,
}
pub struct GeminiClient { pub struct GeminiClient {
api_key: String, api_key: String,
@@ -48,26 +22,37 @@ pub struct GeminiClient {
impl GeminiClient { impl GeminiClient {
pub fn new(api_key: String) -> Self { pub fn new(api_key: String) -> Self {
Self::with_model(api_key, "gemini-3-flash-preview".to_string()) Self::with_model(api_key, DEFAULT_MODEL.to_string())
} }
pub fn with_model(api_key: String, model: String) -> Self { 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 { Self {
api_key, api_key,
client: Client::builder() client,
.timeout(Duration::from_secs(30)) base_url,
.build()
.unwrap_or_default(),
base_url: format!(
"https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent",
model
),
model, model,
timeout: Duration::from_secs(30), timeout,
} }
} }
/// Takes a list of filenames and asks Gemini to categorize them 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( pub async fn organize_files(
&self, &self,
filenames: Vec<String>, filenames: Vec<String>,
@@ -75,97 +60,82 @@ impl GeminiClient {
self.organize_files_with_cache(filenames, None, None).await self.organize_files_with_cache(filenames, None, None).await
} }
/// Takes a list of filenames and asks Gemini to categorize them with caching support
pub async fn organize_files_with_cache( pub async fn organize_files_with_cache(
&self, &self,
filenames: Vec<String>, filenames: Vec<String>,
mut cache: Option<&mut Cache>, mut cache: Option<&mut Cache>,
base_path: Option<&Path>, base_path: Option<&Path>,
) -> Result<OrganizationPlan, GeminiError> { ) -> Result<OrganizationPlan, GeminiError> {
let url = format!("{}?key={}", self.base_url, self.api_key); let url = self.build_url();
// Check cache first if available if let (Some(cache), Some(base_path)) = (cache.as_ref(), base_path) {
if let (Some(cache_ref), Some(base_path)) = (cache.as_ref(), base_path) if let Some(cached_response) = cache.get_cached_response(&filenames, base_path) {
&& let Some(cached_response) = cache_ref.get_cached_response(&filenames, base_path)
{
return Ok(cached_response); return Ok(cached_response);
} }
// 1. Construct the Prompt
let file_list_str = filenames.join(", ");
let prompt = format!(
"I have these files in my Downloads folder: [{}]. \
Categorize them into these folders: 'Images', 'Documents', 'Installers', 'Music', 'Archives', 'Code', 'Misc'. \
Return ONLY a JSON object with this structure: {{ 'files': [ {{ 'filename': 'name', 'category': 'folder' }} ] }}",
file_list_str
);
// 2. Build Request with JSON Mode enforced
let request_body = json!({
"contents": [{
"parts": [{ "text": prompt }]
}],
"generationConfig": {
"response_mime_type": "application/json"
} }
});
// 3. Send with retry logic 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 res = self.send_request_with_retry(&url, &request_body).await?;
let plan = self.parse_categorization_response(res).await?;
// 4. Parse
if res.status().is_success() {
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
.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;
// Deserialize into our temporary response struct
let plan_response: OrganizationPlanResponse = serde_json::from_str(raw_text)?;
// Manually map to the final OrganizationPlan
let plan = OrganizationPlan {
files: plan_response
.files
.into_iter()
.map(|f| FileCategory {
filename: f.filename,
category: f.category,
sub_category: String::new(), // Initialize with empty sub_category
})
.collect(),
};
// Cache the response if cache is available
if let (Some(cache), Some(base_path)) = (cache.as_mut(), base_path) { if let (Some(cache), Some(base_path)) = (cache.as_mut(), base_path) {
cache.cache_response(&filenames, plan.clone(), base_path); cache.cache_response(&filenames, plan.clone(), base_path);
} }
Ok(plan) Ok(plan)
} else {
Err(GeminiError::from_response(res).await)
}
} }
/// Send request with retry logic for retryable errors 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( async fn send_request_with_retry(
&self, &self,
url: &str, url: &str,
request_body: &serde_json::Value, request_body: &serde_json::Value,
) -> Result<reqwest::Response, GeminiError> { ) -> Result<reqwest::Response, GeminiError> {
let mut attempts = 0; let mut attempts = 0;
let max_attempts = 3;
let mut base_delay = Duration::from_secs(2); let mut base_delay = Duration::from_secs(2);
loop { loop {
@@ -179,15 +149,9 @@ impl GeminiClient {
let error = GeminiError::from_response(response).await; let error = GeminiError::from_response(response).await;
if error.is_retryable() && attempts < max_attempts { if error.is_retryable() && attempts < MAX_RETRIES {
let delay = error.retry_delay().unwrap_or(base_delay); let delay = error.retry_delay().unwrap_or(base_delay);
println!( self.print_retry_message(&error, delay, attempts);
"API Error: {}. Retrying in {} seconds (attempt {}/{})",
error,
delay.as_secs(),
attempts,
max_attempts
);
tokio::time::sleep(delay).await; tokio::time::sleep(delay).await;
base_delay *= 2; base_delay *= 2;
continue; continue;
@@ -196,14 +160,8 @@ impl GeminiClient {
return Err(error); return Err(error);
} }
Err(e) => { Err(e) => {
if attempts < max_attempts { if attempts < MAX_RETRIES {
println!( self.print_network_retry(&e, base_delay, attempts);
"Network error: {}. Retrying in {} seconds (attempt {}/{})",
e,
base_delay.as_secs(),
attempts,
max_attempts
);
tokio::time::sleep(base_delay).await; tokio::time::sleep(base_delay).await;
base_delay *= 2; base_delay *= 2;
continue; continue;
@@ -214,24 +172,35 @@ impl GeminiClient {
} }
} }
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( pub async fn get_ai_sub_category(
&self, &self,
filename: &str, filename: &str,
parent_category: &str, parent_category: &str,
content: &str, content: &str,
) -> String { ) -> String {
let url = format!("{}?key={}", self.base_url, self.api_key); let url = self.build_url();
let prompt = PromptBuilder::build_subcategory_prompt(filename, parent_category, content);
let prompt = format!( let request_body = self.build_subcategory_request(&prompt);
"I have a file named '{}' inside the '{}' folder. Here is the first 1000 characters of the content:\n---\n{}\n---\nBased on this, suggest a single short sub-folder name (e.g., 'Invoices', 'Notes', 'Config'). Return ONLY the name of the sub-folder. Do not use markdown or explanations.",
filename, parent_category, content
);
let request_body = json!({
"contents": [{
"parts": [{ "text": prompt }]
}]
});
let res = match self.client.post(&url).json(&request_body).send().await { let res = match self.client.post(&url).json(&request_body).send().await {
Ok(res) => res, Ok(res) => res,
@@ -244,7 +213,25 @@ impl GeminiClient {
} }
}; };
if res.status().is_success() { 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 { let gemini_response: GeminiResponse = match res.json().await {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
@@ -253,26 +240,20 @@ impl GeminiClient {
} }
}; };
let sub_category = gemini_response self.extract_subcategory_from_response(&gemini_response, filename)
.candidates }
.first()
.and_then(|c| c.content.parts.first())
.map(|p| p.text.trim())
.unwrap_or("General")
.to_string();
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() { if sub_category.is_empty() {
"General".to_string() "General".to_string()
} else { } else {
sub_category sub_category.to_string()
} }
} else { }
eprintln!( Err(_) => "General".to_string(),
"Warning: API returned error for {}: {}",
filename,
res.status()
);
"General".to_string()
} }
} }
} }

51
src/gemini_helpers.rs Normal file
View File

@@ -0,0 +1,51 @@
use crate::files::{FileCategory, OrganizationPlan};
use crate::gemini_types::OrganizationPlanResponse;
impl OrganizationPlanResponse {
pub fn to_organization_plan(self) -> OrganizationPlan {
OrganizationPlan {
files: self
.files
.into_iter()
.map(|f| FileCategory {
filename: f.filename,
category: f.category,
sub_category: String::new(),
})
.collect(),
}
}
}
#[derive(Debug)]
pub struct PromptBuilder {
file_list: String,
}
impl PromptBuilder {
pub fn new(file_list: Vec<String>) -> Self {
Self {
file_list: file_list.join(", "),
}
}
pub fn build_categorization_prompt(&self) -> String {
format!(
"I have these files in my Downloads folder: [{}]. \
Categorize them into these folders: 'Images', 'Documents', 'Installers', 'Music', 'Archives', 'Code', 'Misc'. \
Return ONLY a JSON object with this structure: {{ 'files': [ {{ 'filename': 'name', 'category': 'folder' }} ] }}",
self.file_list
)
}
pub fn build_subcategory_prompt(
filename: &str,
parent_category: &str,
content: &str,
) -> String {
format!(
"I have a file named '{}' inside the '{}' folder. Here is the first 1000 characters of content:\n---\n{}\n---\nBased on this, suggest a single short sub-folder name (e.g., 'Invoices', 'Notes', 'Config'). Return ONLY the name of the sub-folder. Do not use markdown or explanations.",
filename, parent_category, content
)
}
}

32
src/gemini_types.rs Normal file
View File

@@ -0,0 +1,32 @@
use serde::Deserialize;
#[derive(Deserialize, Default)]
pub struct GeminiResponse {
pub candidates: Vec<Candidate>,
}
#[derive(Deserialize)]
pub struct Candidate {
pub content: Content,
}
#[derive(Deserialize)]
pub struct Content {
pub parts: Vec<Part>,
}
#[derive(Deserialize)]
pub struct Part {
pub text: String,
}
#[derive(Deserialize)]
pub struct FileCategoryResponse {
pub filename: String,
pub category: String,
}
#[derive(Deserialize)]
pub struct OrganizationPlanResponse {
pub files: Vec<FileCategoryResponse>,
}

View File

@@ -3,3 +3,6 @@ pub mod config;
pub mod files; pub mod files;
pub mod gemini; pub mod gemini;
pub mod gemini_errors; pub mod gemini_errors;
pub mod gemini_helpers;
pub mod gemini_types;
pub mod prompt;

130
src/prompt.rs Normal file
View File

@@ -0,0 +1,130 @@
use colored::*;
use directories::BaseDirs;
use std::io::Write;
use std::path::{Path, PathBuf};
const MAX_RETRIES: u32 = 3;
pub struct Prompter;
impl Prompter {
pub fn prompt_api_key() -> Result<String, Box<dyn std::error::Error>> {
println!();
println!("Get your API key at: {}", "https://ai.google.dev/".cyan().underline());
println!("Enter your API Key (starts with 'AIza'):");
let mut attempts = 0;
while attempts < MAX_RETRIES {
print!("API Key: ");
std::io::stdout().flush()?;
let mut input = String::new();
std::io::stdin().read_line(&mut input)?;
let key = input.trim();
if Self::validate_api_key(key) {
return Ok(key.to_string());
}
attempts += 1;
Self::print_validation_error("Invalid API key format. Must start with 'AIza' and be around 39 characters.", attempts);
}
Err("Max retries exceeded. Please run again with a valid API key.".into())
}
pub fn prompt_download_folder() -> Result<PathBuf, Box<dyn std::error::Error>> {
let default_path = Self::get_default_downloads_folder();
let default_display = default_path.to_string_lossy();
println!();
println!(
"Enter path to folder to organize (e.g., {}):",
default_display.yellow()
);
println!("Or press Enter to use default: {}", default_display.green());
println!("Folder path: ");
let mut attempts = 0;
while attempts < MAX_RETRIES {
print!("> ");
std::io::stdout().flush()?;
let mut input = String::new();
std::io::stdin().read_line(&mut input)?;
let input = input.trim();
let path = if input.is_empty() {
default_path.clone()
} else {
let expanded = Self::expand_home(input);
PathBuf::from(expanded)
};
if Self::validate_folder_path(&path) {
return Ok(path);
}
attempts += 1;
Self::print_path_validation_error(&path, attempts);
}
Err("Max retries exceeded. Please run again with a valid folder path.".into())
}
fn print_validation_error(message: &str, attempts: u32) {
let remaining = MAX_RETRIES - attempts;
eprintln!("{} {}", "".red(), message);
if remaining > 0 {
eprintln!("Try again ({} attempts remaining):", remaining);
}
}
fn print_path_validation_error(path: &Path, attempts: u32) {
let remaining = MAX_RETRIES - attempts;
eprintln!("{} Invalid folder path.", "".red());
if !path.exists() {
eprintln!(" Path does not exist: {}", path.display());
} else if !path.is_dir() {
eprintln!(" Path is not a directory: {}", path.display());
}
if remaining > 0 {
eprintln!("Try again ({} attempts remaining):", remaining);
println!("Folder path: ");
}
}
pub fn validate_api_key(key: &str) -> bool {
!key.is_empty()
&& key.starts_with("AIza")
&& key.len() >= 35
&& key.len() <= 50
}
pub fn validate_folder_path(path: &Path) -> bool {
path.exists() && path.is_dir()
}
pub fn get_default_downloads_folder() -> PathBuf {
BaseDirs::new()
.map(|base_dirs| base_dirs.home_dir().join("Downloads"))
.unwrap_or_else(|| PathBuf::from("./Downloads"))
}
pub fn expand_home(path: &str) -> String {
if path.starts_with("~/") {
if let Some(base_dirs) = BaseDirs::new() {
let home = base_dirs.home_dir();
return path.replacen("~", &home.to_string_lossy(), 1);
}
}
path.to_string()
}
}