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:
204
src/config.rs
204
src/config.rs
@@ -1,11 +1,6 @@
|
||||
use colored::*;
|
||||
use directories::{BaseDirs, ProjectDirs};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fs;
|
||||
use std::io::{self, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
const MAX_RETRIES: u32 = 3;
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct Config {
|
||||
@@ -14,28 +9,14 @@ pub struct Config {
|
||||
}
|
||||
|
||||
impl Config {
|
||||
fn get_config_dir() -> Result<PathBuf, 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>> {
|
||||
pub fn load() -> Result<Self, Box<dyn std::error::Error>> {
|
||||
let config_path = Self::get_config_path()?;
|
||||
|
||||
if !config_path.exists() {
|
||||
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)?;
|
||||
|
||||
if config.api_key.is_empty() {
|
||||
@@ -50,7 +31,7 @@ impl Config {
|
||||
|
||||
let toml_string = toml::to_string_pretty(self)?;
|
||||
|
||||
fs::write(&config_path, toml_string)?;
|
||||
std::fs::write(&config_path, toml_string)?;
|
||||
|
||||
println!(
|
||||
"{} Configuration saved to {}",
|
||||
@@ -74,6 +55,21 @@ impl Config {
|
||||
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>> {
|
||||
@@ -87,17 +83,9 @@ pub fn get_or_prompt_api_key() -> Result<String, Box<dyn std::error::Error>> {
|
||||
println!("{}", "🔑 NoEntropy Configuration".bold().cyan());
|
||||
println!("{}", "─────────────────────────────".cyan());
|
||||
|
||||
let api_key = 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 api_key = crate::prompt::Prompter::prompt_api_key()?;
|
||||
|
||||
let mut config = Config::load().unwrap_or_default();
|
||||
config.api_key = api_key.clone();
|
||||
config.save()?;
|
||||
|
||||
@@ -115,17 +103,9 @@ pub fn get_or_prompt_download_folder() -> Result<PathBuf, Box<dyn std::error::Er
|
||||
println!();
|
||||
println!("{}", "📁 Download folder not configured.".yellow());
|
||||
|
||||
let folder_path = 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 folder_path = crate::prompt::Prompter::prompt_download_folder()?;
|
||||
|
||||
let mut config = Config::load().unwrap_or_default();
|
||||
config.download_folder = folder_path.clone();
|
||||
config.save()?;
|
||||
|
||||
@@ -133,141 +113,13 @@ pub fn get_or_prompt_download_folder() -> Result<PathBuf, Box<dyn std::error::Er
|
||||
Ok(folder_path)
|
||||
}
|
||||
|
||||
fn prompt_api_key() -> Result<String, Box<dyn std::error::Error>> {
|
||||
let mut attempts = 0;
|
||||
|
||||
println!();
|
||||
println!("Get your API key at: {}", "https://ai.google.dev/".cyan().underline());
|
||||
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);
|
||||
impl Default for Config {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
api_key: String::new(),
|
||||
download_folder: PathBuf::new(),
|
||||
}
|
||||
}
|
||||
|
||||
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)]
|
||||
|
||||
2
src/config_mod.rs.bak
Normal file
2
src/config_mod.rs.bak
Normal file
@@ -0,0 +1,2 @@
|
||||
pub mod config;
|
||||
pub mod prompt;
|
||||
@@ -1,4 +1,5 @@
|
||||
use crate::config::*;
|
||||
use std::path::Path;
|
||||
|
||||
#[test]
|
||||
fn test_config_serialization() {
|
||||
@@ -17,37 +18,37 @@ fn test_config_serialization() {
|
||||
|
||||
#[test]
|
||||
fn test_validate_api_key_valid() {
|
||||
assert!(validate_api_key("AIzaSyB1234567890123456789012345678"));
|
||||
assert!(validate_api_key("AIzaSyB123456789012345678901234567890"));
|
||||
assert!(crate::prompt::Prompter::validate_api_key("AIzaSyB1234567890123456789012345678"));
|
||||
assert!(crate::prompt::Prompter::validate_api_key("AIzaSyB123456789012345678901234567890"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validate_api_key_invalid() {
|
||||
assert!(!validate_api_key(""));
|
||||
assert!(!validate_api_key("invalid_key"));
|
||||
assert!(!validate_api_key("BizaSyB1234567890123456789012345678"));
|
||||
assert!(!validate_api_key("short"));
|
||||
assert!(!crate::prompt::Prompter::validate_api_key(""));
|
||||
assert!(!crate::prompt::Prompter::validate_api_key("invalid_key"));
|
||||
assert!(!crate::prompt::Prompter::validate_api_key("BizaSyB1234567890123456789012345678"));
|
||||
assert!(!crate::prompt::Prompter::validate_api_key("short"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validate_folder_path_valid() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
assert!(validate_folder_path(temp_dir.path()));
|
||||
assert!(crate::prompt::Prompter::validate_folder_path(temp_dir.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
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();
|
||||
assert!(!validate_folder_path(temp_file.path()));
|
||||
assert!(!crate::prompt::Prompter::validate_folder_path(temp_file.path()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
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 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.contains("test/path"));
|
||||
}
|
||||
@@ -55,16 +56,16 @@ fn test_expand_home_with_tilde() {
|
||||
|
||||
#[test]
|
||||
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");
|
||||
|
||||
let expanded = expand_home("relative/path");
|
||||
let expanded = crate::prompt::Prompter::expand_home("relative/path");
|
||||
assert_eq!(expanded, "relative/path");
|
||||
}
|
||||
|
||||
#[test]
|
||||
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"));
|
||||
}
|
||||
|
||||
|
||||
307
src/gemini.rs
307
src/gemini.rs
@@ -1,42 +1,16 @@
|
||||
use crate::cache::Cache;
|
||||
use crate::files::{FileCategory, OrganizationPlan};
|
||||
use crate::gemini_errors::GeminiError;
|
||||
use crate::gemini_helpers::PromptBuilder;
|
||||
use crate::gemini_types::{GeminiResponse, OrganizationPlanResponse};
|
||||
use crate::files::OrganizationPlan;
|
||||
use reqwest::Client;
|
||||
use serde::Deserialize;
|
||||
use serde_json::json;
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
#[derive(Deserialize, Default)]
|
||||
struct GeminiResponse {
|
||||
candidates: Vec<Candidate>,
|
||||
}
|
||||
|
||||
#[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>,
|
||||
}
|
||||
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,
|
||||
@@ -48,26 +22,37 @@ pub struct GeminiClient {
|
||||
|
||||
impl GeminiClient {
|
||||
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 {
|
||||
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: Client::builder()
|
||||
.timeout(Duration::from_secs(30))
|
||||
.build()
|
||||
.unwrap_or_default(),
|
||||
base_url: format!(
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent",
|
||||
model
|
||||
),
|
||||
client,
|
||||
base_url,
|
||||
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(
|
||||
&self,
|
||||
filenames: Vec<String>,
|
||||
@@ -75,97 +60,82 @@ impl GeminiClient {
|
||||
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(
|
||||
&self,
|
||||
filenames: Vec<String>,
|
||||
mut cache: Option<&mut Cache>,
|
||||
base_path: Option<&Path>,
|
||||
) -> 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_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);
|
||||
if let (Some(cache), Some(base_path)) = (cache.as_ref(), base_path) {
|
||||
if let Some(cached_response) = cache.get_cached_response(&filenames, base_path) {
|
||||
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
|
||||
);
|
||||
let prompt = PromptBuilder::new(filenames.clone()).build_categorization_prompt();
|
||||
let request_body = self.build_categorization_request(&prompt);
|
||||
|
||||
// 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 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) {
|
||||
cache.cache_response(&filenames, plan.clone(), base_path);
|
||||
}
|
||||
|
||||
Ok(plan)
|
||||
} else {
|
||||
Err(GeminiError::from_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())
|
||||
}
|
||||
|
||||
/// Send request with retry logic for retryable errors
|
||||
async fn send_request_with_retry(
|
||||
&self,
|
||||
url: &str,
|
||||
request_body: &serde_json::Value,
|
||||
) -> Result<reqwest::Response, GeminiError> {
|
||||
let mut attempts = 0;
|
||||
let max_attempts = 3;
|
||||
let mut base_delay = Duration::from_secs(2);
|
||||
|
||||
loop {
|
||||
@@ -179,15 +149,9 @@ impl GeminiClient {
|
||||
|
||||
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);
|
||||
println!(
|
||||
"API Error: {}. Retrying in {} seconds (attempt {}/{})",
|
||||
error,
|
||||
delay.as_secs(),
|
||||
attempts,
|
||||
max_attempts
|
||||
);
|
||||
self.print_retry_message(&error, delay, attempts);
|
||||
tokio::time::sleep(delay).await;
|
||||
base_delay *= 2;
|
||||
continue;
|
||||
@@ -196,14 +160,8 @@ impl GeminiClient {
|
||||
return Err(error);
|
||||
}
|
||||
Err(e) => {
|
||||
if attempts < max_attempts {
|
||||
println!(
|
||||
"Network error: {}. Retrying in {} seconds (attempt {}/{})",
|
||||
e,
|
||||
base_delay.as_secs(),
|
||||
attempts,
|
||||
max_attempts
|
||||
);
|
||||
if attempts < MAX_RETRIES {
|
||||
self.print_network_retry(&e, base_delay, attempts);
|
||||
tokio::time::sleep(base_delay).await;
|
||||
base_delay *= 2;
|
||||
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(
|
||||
&self,
|
||||
filename: &str,
|
||||
parent_category: &str,
|
||||
content: &str,
|
||||
) -> String {
|
||||
let url = format!("{}?key={}", self.base_url, self.api_key);
|
||||
|
||||
let prompt = format!(
|
||||
"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 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,
|
||||
@@ -244,35 +213,47 @@ impl GeminiClient {
|
||||
}
|
||||
};
|
||||
|
||||
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();
|
||||
}
|
||||
};
|
||||
self.parse_subcategory_response(res, filename).await
|
||||
}
|
||||
|
||||
let sub_category = gemini_response
|
||||
.candidates
|
||||
.first()
|
||||
.and_then(|c| c.content.parts.first())
|
||||
.map(|p| p.text.trim())
|
||||
.unwrap_or("General")
|
||||
.to_string();
|
||||
fn build_subcategory_request(&self, prompt: &str) -> serde_json::Value {
|
||||
json!({
|
||||
"contents": [{ "parts": [{ "text": prompt }] }]
|
||||
})
|
||||
}
|
||||
|
||||
if sub_category.is_empty() {
|
||||
"General".to_string()
|
||||
} else {
|
||||
sub_category
|
||||
}
|
||||
} else {
|
||||
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()
|
||||
);
|
||||
"General".to_string()
|
||||
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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
51
src/gemini_helpers.rs
Normal file
51
src/gemini_helpers.rs
Normal 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
32
src/gemini_types.rs
Normal 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>,
|
||||
}
|
||||
@@ -3,3 +3,6 @@ pub mod config;
|
||||
pub mod files;
|
||||
pub mod gemini;
|
||||
pub mod gemini_errors;
|
||||
pub mod gemini_helpers;
|
||||
pub mod gemini_types;
|
||||
pub mod prompt;
|
||||
|
||||
130
src/prompt.rs
Normal file
130
src/prompt.rs
Normal 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()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user