siddheshmhatre/noentropy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: Simplify codebase by extracting modules and helpers

Browse Source 
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:
glitchySid
2025-12-29 00:35:14 +05:30
parent 32c1f283ac
commit 1a72116b9d
8 changed files with 405 additions and 353 deletions
Download Patch File Download Diff File Expand all files Collapse all files

307
src/gemini.rs
View File

@@ -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(),
}
}
}