use crate::cache::Cache; use crate::files::{FileCategory, OrganizationPlan}; use crate::gemini_errors::GeminiError; 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, } #[derive(Deserialize)] struct Candidate { content: Content, } #[derive(Deserialize)] struct Content { parts: Vec, } #[derive(Deserialize)] struct Part { text: String, } #[derive(Deserialize)] struct FileCategoryResponse { filename: String, category: String, } #[derive(Deserialize)] struct OrganizationPlanResponse { files: Vec, } pub struct GeminiClient { api_key: String, client: Client, base_url: String, } impl GeminiClient { pub fn new(api_key: String) -> Self { Self { api_key, client: Client::new(), base_url: "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent".to_string(), } } /// Takes a list of filenames and asks Gemini to categorize them pub async fn organize_files( &self, filenames: Vec, ) -> Result { 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, mut cache: Option<&mut Cache>, base_path: Option<&Path>, ) -> Result { let url = format!("{}?key={}", self.base_url, self.api_key); // Check cache first if available if let (Some(cache_ref), Some(base_path)) = (cache.as_ref(), base_path) { if let Some(cached_response) = cache_ref.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 ); // 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?; // 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 .get(0) .ok_or_else(|| GeminiError::InvalidResponse("No candidates in response".to_string()))? .content.parts .get(0) .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) } } /// Send request with retry logic for retryable errors async fn send_request_with_retry( &self, url: &str, request_body: &serde_json::Value, ) -> Result { let mut attempts = 0; let max_attempts = 3; loop { attempts += 1; match self.client.post(url).json(request_body).send().await { Ok(response) => { if response.status().is_success() { return Ok(response); } let error = GeminiError::from_response(response).await; if error.is_retryable() && attempts < max_attempts { if let Some(delay) = error.retry_delay() { println!("API Error: {}. Retrying in {} seconds (attempt {}/{})", error, delay.as_secs(), attempts, max_attempts); tokio::time::sleep(delay).await; continue; } } return Err(error); } Err(e) => { if attempts < max_attempts { println!("Network error: {}. Retrying in {} seconds (attempt {}/{})", e, 5, attempts, max_attempts); tokio::time::sleep(Duration::from_secs(5)).await; continue; } return Err(GeminiError::NetworkError(e)); } } } } 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 res = self.client.post(&url).json(&request_body).send().await; if let Ok(res) = res { if res.status().is_success() { let gemini_response: GeminiResponse = res.json().await.unwrap_or_default(); let sub_category = gemini_response.candidates .get(0) .and_then(|c| c.content.parts.get(0)) .map(|p| p.text.trim()) .unwrap_or("General") .to_string(); if sub_category.is_empty() { "General".to_string() } else { sub_category } } else { "General".to_string() } } else { "General".to_string() } } }