Merge pull request #8 from glitchySid/feature/customcategories

Feature/customcategories

Cargo.lock

@@ -721,7 +721,7 @@ dependencies = [
 
 [[package]]
 name = "noentropy"
-version = "1.0.3"
+version = "1.0.4"
 dependencies = [
  "clap",
  "colored",

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "noentropy"
-version = "1.0.3"
+version = "1.0.4"
 edition = "2024"
 
 [dependencies]

README.md

@@ -21,6 +21,7 @@ NoEntropy is a smart command-line tool that organizes your cluttered Downloads f
 ## Features
 
 - **🧠 AI-Powered Categorization** - Uses Google Gemini API for intelligent file sorting
+- **🎨 Custom Categories** - Define your own categories for personalized organization
 - **📁 Automatic Sub-Folders** - Creates relevant sub-folders based on file content analysis
 - **💨 Smart Caching** - Minimizes API calls with metadata-based caching (7-day expiry)
 - **⚡ Concurrent Processing** - Parallel file inspection with configurable limits
@@ -83,12 +84,16 @@ NoEntropy stores configuration in `~/.config/noentropy/config.toml` following XD
 ```toml
 api_key = "AIzaSyDTEhAq414SHY094A5oy5lxNA0vhbY1O3k"
 download_folder = "/home/user/Downloads"
+
+# Optional: Custom categories for file organization
+categories = ["Work", "Personal", "School", "Projects", "Bills", "Media", "Misc"]
 ```
 
-| Setting | Description | Example |
+| Setting | Description | Example | Required |
-|---------|-------------|---------|
+|---------|-------------|---------|----------|
-| `api_key` | Your Google Gemini API key | `AIzaSy...` |
+| `api_key` | Your Google Gemini API key | `AIzaSy...` | Yes |
-| `download_folder` | Path to folder to organize | `/home/user/Downloads` |
+| `download_folder` | Path to folder to organize | `/home/user/Downloads` | Yes |
+| `categories` | Custom categories for organization | `["Work", "Personal", "School"]` | No |
 
 ### Getting a Gemini API Key
 
@@ -107,6 +112,89 @@ NoEntropy provides an interactive setup on first run:
 
 Configuration is automatically saved to `~/.config/noentropy/config.toml` after interactive setup.
 
+### Custom Categories
+
+NoEntropy allows you to define your own custom categories instead of using the default ones. This is perfect for organizing files based on your specific workflow or needs.
+
+#### Default Categories
+
+If you don't specify custom categories, NoEntropy uses these defaults:
+- **Images** - PNG, JPG, GIF, SVG, etc.
+- **Documents** - PDF, DOC, DOCX, TXT, MD, etc.
+- **Installers** - EXE, DMG, APP, PKG, etc.
+- **Music** - MP3, WAV, FLAC, M4A, etc.
+- **Archives** - ZIP, TAR, RAR, 7Z, etc.
+- **Code** - Source code and configuration files
+- **Misc** - Everything else
+
+#### Using Custom Categories
+
+To use custom categories, add a `categories` array to your `config.toml`:
+
+```toml
+api_key = "your_api_key_here"
+download_folder = "/home/user/Downloads"
+categories = ["Work", "Personal", "School", "Projects", "Bills", "Media", "Misc"]
+```
+
+**Examples of Custom Category Sets:**
+
+**For Students:**
+```toml
+categories = ["Courses", "Assignments", "Research", "Personal", "Textbooks", "Media", "Misc"]
+```
+
+**For Professionals:**
+```toml
+categories = ["Client Work", "Internal", "Invoices", "Contracts", "Marketing", "Resources", "Misc"]
+```
+
+**For Creatives:**
+```toml
+categories = ["Projects", "Assets", "References", "Client Files", "Portfolio", "Tools", "Misc"]
+```
+
+**For Personal Use:**
+```toml
+categories = ["Family", "Finance", "Health", "Home", "Travel", "Hobbies", "Misc"]
+```
+
+#### Tips for Custom Categories
+
+1. **Keep it simple** - Use 5-10 categories for best results
+2. **Be specific** - More descriptive names help the AI understand better
+3. **Include "Misc"** - Always have a catch-all category for unclear files
+4. **Think workflow** - Organize based on how you actually use files
+5. **Test first** - Use `--dry-run` to preview categorization before committing
+
+#### How It Works
+
+When you define custom categories:
+1. NoEntropy sends your file list to the Gemini AI
+2. The AI is instructed to categorize files into your custom categories
+3. Files are organized into folders matching your category names
+4. Sub-folders are still created automatically for better organization
+
+**Example Output with Custom Categories:**
+```
+Downloads/
+├── Work/
+│   ├── Reports/
+│   │   └── Q4-Report.pdf
+│   └── Presentations/
+│       └── Client-Deck.pptx
+├── Personal/
+│   ├── Photos/
+│   │   └── vacation.jpg
+│   └── Documents/
+│       └── resume.pdf
+└── School/
+    ├── Assignments/
+    │   └── homework.docx
+    └── Notes/
+        └── lecture-notes.pdf
+```
+
 ## Usage
 
 ### Basic Usage
@@ -498,7 +586,7 @@ noentropy/
 
 Based on community feedback, we're planning:
 
-- [ ] **Custom Categories** - Define custom categories in `config.toml`
+- [x] **Custom Categories** - Define custom categories in `config.toml`
 - [x] **Recursive Mode** - Organize files in subdirectories with `--recursive` flag
 - [x] **Undo Functionality** - Revert file organization changes
 - [ ] **Custom Models** - Support for other AI providers

config.example.toml

@@ -7,3 +7,8 @@ api_key = "your_api_key_here"
 
 # Path to folder to organize (e.g., ~/Downloads)
 download_folder = "/path/to/your/downloads"
+
+# Optional: Custom categories for file organization
+# If not specified, uses default categories: Images, Documents, Installers, Music, Archives, Code, Misc
+# Uncomment and customize the list below to use your own categories:
+# categories = ["Work", "Personal", "School", "Projects", "Bills", "Media", "Misc"]

src/cli/orchestrator.rs

@@ -91,10 +91,9 @@ pub fn handle_gemini_error(error: crate::gemini::GeminiError) {
 
 pub async fn handle_organization(
     args: Args,
-    api_key: String,
+    config: Config,
-    download_path: PathBuf,
 ) -> Result<(), Box<dyn std::error::Error>> {
-    let client: GeminiClient = GeminiClient::new(api_key);
+    let client: GeminiClient = GeminiClient::new(config.api_key, config.categories.clone());
 
     let data_dir = Config::get_data_dir()?;
     let cache_path = data_dir.join(".noentropy_cache.json");
@@ -106,6 +105,7 @@ pub async fn handle_organization(
     let mut undo_log = UndoLog::load_or_create(&undo_log_path);
     undo_log.cleanup_old_entries(30 * 24 * 60 * 60);
 
+    let download_path = config.download_folder;
     let batch = FileBatch::from_path(download_path.clone(), args.recursive);
 
     if batch.filenames.is_empty() {

src/gemini/client.rs

@@ -21,14 +21,15 @@ pub struct GeminiClient {
     model: String,
     #[allow(dead_code)]
     timeout: Duration,
+    categories: Vec<String>,
 }
 
 impl GeminiClient {
-    pub fn new(api_key: String) -> Self {
+    pub fn new(api_key: String, categories: Vec<String>) -> Self {
-        Self::with_model(api_key, DEFAULT_MODEL.to_string())
+        Self::with_model(api_key, DEFAULT_MODEL.to_string(), categories)
     }
 
-    pub fn with_model(api_key: String, model: String) -> Self {
+    pub fn with_model(api_key: String, model: String, categories: Vec<String>) -> Self {
         let timeout = Duration::from_secs(DEFAULT_TIMEOUT_SECS);
         let client = Self::build_client(timeout);
         let base_url = Self::build_base_url(&model);
@@ -39,6 +40,7 @@ impl GeminiClient {
             base_url,
             model,
             timeout,
+            categories,
         }
     }
 
@@ -77,7 +79,8 @@ impl GeminiClient {
             return Ok(cached_response);
         }
 
-        let prompt = PromptBuilder::new(filenames.clone()).build_categorization_prompt();
+        let prompt =
+            PromptBuilder::new(filenames.clone()).build_categorization_prompt(&self.categories);
         let request_body = self.build_categorization_request(&prompt);
 
         let res = self.send_request_with_retry(&url, &request_body).await?;

src/gemini/prompt.rs

@@ -29,12 +29,13 @@ impl PromptBuilder {
         }
     }
 
-    pub fn build_categorization_prompt(&self) -> String {
+    pub fn build_categorization_prompt(&self, categories: &[String]) -> String {
+        let categories_str = categories.join("', '");
         format!(
             "I have these files in my Downloads folder: [{}]. \
-             Categorize them into these folders: 'Images', 'Documents', 'Installers', 'Music', 'Archives', 'Code', 'Misc'. \
+             Categorize them into these folders: '{}'. \
              Return ONLY a JSON object with this structure: {{ 'files': [ {{ 'filename': 'name', 'category': 'folder' }} ] }}",
-            self.file_list
+            self.file_list, categories_str
         )
     }
 

src/main.rs

@@ -3,7 +3,7 @@ use noentropy::cli::{
     Args,
     orchestrator::{handle_organization, handle_undo},
 };
-use noentropy::settings::{get_or_prompt_api_key, get_or_prompt_download_folder};
+use noentropy::settings::{get_or_prompt_config, get_or_prompt_download_folder};
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
@@ -15,10 +15,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         return Ok(());
     }
 
-    let api_key = get_or_prompt_api_key()?;
+    let config = get_or_prompt_config()?;
-    let download_path = get_or_prompt_download_folder()?;
 
-    handle_organization(args, api_key, download_path).await?;
+    handle_organization(args, config).await?;
 
     Ok(())
 }

src/settings/config.rs

@@ -5,10 +5,24 @@ use std::path::PathBuf;
 
 use super::prompt::Prompter;
 
+pub fn default_categories() -> Vec<String> {
+    vec![
+        "Images".to_string(),
+        "Documents".to_string(),
+        "Installers".to_string(),
+        "Music".to_string(),
+        "Archives".to_string(),
+        "Code".to_string(),
+        "Misc".to_string(),
+    ]
+}
+
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct Config {
     pub api_key: String,
     pub download_folder: PathBuf,
+    #[serde(default = "default_categories")]
+    pub categories: Vec<String>,
 }
 
 impl Config {
@@ -128,11 +142,41 @@ pub fn get_or_prompt_download_folder() -> Result<PathBuf, Box<dyn std::error::Er
     Ok(folder_path)
 }
 
+pub fn get_or_prompt_config() -> Result<Config, Box<dyn std::error::Error>> {
+    let mut config = Config::load().unwrap_or_default();
+    let mut needs_save = false;
+
+    // Check API key
+    if config.api_key.is_empty() {
+        println!();
+        println!("{}", "🔑 NoEntropy Configuration".bold().cyan());
+        println!("{}", "─────────────────────────────".cyan());
+        config.api_key = Prompter::prompt_api_key()?;
+        needs_save = true;
+    }
+
+    // Check download folder
+    if config.download_folder.as_os_str().is_empty() || !config.download_folder.exists() {
+        println!();
+        println!("{}", "📁 Download folder not configured.".yellow());
+        config.download_folder = Prompter::prompt_download_folder()?;
+        needs_save = true;
+    }
+
+    if needs_save {
+        config.save()?;
+        println!();
+    }
+
+    Ok(config)
+}
+
 impl Default for Config {
     fn default() -> Self {
         Self {
             api_key: String::new(),
             download_folder: PathBuf::new(),
+            categories: default_categories(),
         }
     }
 }

src/settings/mod.rs

@@ -1,7 +1,9 @@
 pub mod config;
 pub mod prompt;
 
-pub use config::{Config, get_or_prompt_api_key, get_or_prompt_download_folder};
+pub use config::{
+    Config, get_or_prompt_api_key, get_or_prompt_config, get_or_prompt_download_folder,
+};
 pub use prompt::Prompter;
 
 #[cfg(test)]

src/settings/tests.rs

@@ -1,3 +1,4 @@
+use crate::settings::config::default_categories;
 use crate::settings::*;
 use std::path::Path;
 use std::path::PathBuf;
@@ -7,6 +8,7 @@ fn test_config_serialization() {
     let config = Config {
         api_key: "test_key_12345".to_string(),
         download_folder: PathBuf::from("/test/path"),
+        categories: default_categories(),
     };
 
     let toml_str = toml::to_string_pretty(&config).unwrap();
@@ -15,6 +17,7 @@ fn test_config_serialization() {
     let deserialized: Config = toml::from_str(&toml_str).unwrap();
     assert_eq!(config.api_key, deserialized.api_key);
     assert_eq!(config.download_folder, deserialized.download_folder);
+    assert_eq!(config.categories, deserialized.categories);
 }
 
 #[test]
@@ -83,6 +86,7 @@ fn test_config_empty_api_key_error() {
     let config = Config {
         api_key: String::new(),
         download_folder: PathBuf::from("/test/path"),
+        categories: default_categories(),
     };
 
     assert!(config.api_key.is_empty());