Add custom categories feature

- Add support for user-defined custom categories in config.toml - Update Config struct with categories field and default_categories() function - Thread categories through GeminiClient and prompt builder - Update AI prompts to use dynamic categories instead of hardcoded ones - Add comprehensive documentation with examples for different use cases - Update tests to support new categories field - Maintain backward compatibility with default categories - Update version from 1.0.3 to 1.0.4 Closes feature request for custom categories.
2025-12-31 13:01:15 +05:30
parent 08a272c4de
commit 1f0547a210
10 changed files with 167 additions and 21 deletions
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "noentropy"
-version = "1.0.3"
+version = "1.0.4"
 edition = "2024"

 [dependencies]
--- a/README.md
+++ b/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 |
-|---------|-------------|---------|
-| `api_key` | Your Google Gemini API key | `AIzaSy...` |
-| `download_folder` | Path to folder to organize | `/home/user/Downloads` |
+| Setting | Description | Example | Required |
+|---------|-------------|---------|----------|
+| `api_key` | Your Google Gemini API key | `AIzaSy...` | Yes |
+| `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
--- a/config.example.toml
+++ b/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"]
--- a/src/cli/orchestrator.rs
+++ b/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,
-    download_path: PathBuf,
+    config: Config,
 ) -> 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() {
--- a/src/gemini/client.rs
+++ b/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 {
-        Self::with_model(api_key, DEFAULT_MODEL.to_string())
+    pub fn new(api_key: String, categories: Vec<String>) -> Self {
+        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?;
--- a/src/gemini/prompt.rs
+++ b/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
        )
    }

--- a/src/main.rs
+++ b/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 download_path = get_or_prompt_download_folder()?;
+    let config = get_or_prompt_config()?;

-    handle_organization(args, api_key, download_path).await?;
+    handle_organization(args, config).await?;

    Ok(())
 }
--- a/src/settings/config.rs
+++ b/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(),
        }
    }
 }
--- a/src/settings/mod.rs
+++ b/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)]
--- a/src/settings/tests.rs
+++ b/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());