refactored every component thoroughly.
This commit is contained in:
262
src/gemini/client.rs
Normal file
262
src/gemini/client.rs
Normal file
@@ -0,0 +1,262 @@
|
||||
use crate::models::OrganizationPlan;
|
||||
use crate::storage::Cache;
|
||||
use crate::gemini::errors::GeminiError;
|
||||
use crate::gemini::prompt::PromptBuilder;
|
||||
use crate::gemini::types::{GeminiResponse, OrganizationPlanResponse};
|
||||
use reqwest::Client;
|
||||
use serde_json::json;
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
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,
|
||||
client: Client,
|
||||
base_url: String,
|
||||
#[allow(dead_code)]
|
||||
model: String,
|
||||
#[allow(dead_code)]
|
||||
timeout: Duration,
|
||||
}
|
||||
|
||||
impl GeminiClient {
|
||||
pub fn new(api_key: String) -> Self {
|
||||
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,
|
||||
base_url,
|
||||
model,
|
||||
timeout,
|
||||
}
|
||||
}
|
||||
|
||||
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>,
|
||||
) -> Result<OrganizationPlan, GeminiError> {
|
||||
self.organize_files_with_cache(filenames, None, None).await
|
||||
}
|
||||
|
||||
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 = self.build_url();
|
||||
|
||||
if let (Some(cache), Some(base_path)) = (cache.as_ref(), base_path)
|
||||
&& let Some(cached_response) = cache.get_cached_response(&filenames, base_path)
|
||||
{
|
||||
return Ok(cached_response);
|
||||
}
|
||||
|
||||
let prompt = PromptBuilder::new(filenames.clone()).build_categorization_prompt();
|
||||
let request_body = self.build_categorization_request(&prompt);
|
||||
|
||||
let res = self.send_request_with_retry(&url, &request_body).await?;
|
||||
let plan = self.parse_categorization_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())
|
||||
}
|
||||
|
||||
async fn send_request_with_retry(
|
||||
&self,
|
||||
url: &str,
|
||||
request_body: &serde_json::Value,
|
||||
) -> Result<reqwest::Response, GeminiError> {
|
||||
let mut attempts = 0;
|
||||
let mut base_delay = Duration::from_secs(2);
|
||||
|
||||
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_RETRIES {
|
||||
let delay = error.retry_delay().unwrap_or(base_delay);
|
||||
self.print_retry_message(&error, delay, attempts);
|
||||
tokio::time::sleep(delay).await;
|
||||
base_delay *= 2;
|
||||
continue;
|
||||
}
|
||||
|
||||
return Err(error);
|
||||
}
|
||||
Err(e) => {
|
||||
if attempts < MAX_RETRIES {
|
||||
self.print_network_retry(&e, base_delay, attempts);
|
||||
tokio::time::sleep(base_delay).await;
|
||||
base_delay *= 2;
|
||||
continue;
|
||||
}
|
||||
return Err(GeminiError::NetworkError(e));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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 = 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,
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"Warning: Failed to get sub-category for {}: {}",
|
||||
filename, e
|
||||
);
|
||||
return "General".to_string();
|
||||
}
|
||||
};
|
||||
|
||||
self.parse_subcategory_response(res, filename).await
|
||||
}
|
||||
|
||||
fn build_subcategory_request(&self, prompt: &str) -> serde_json::Value {
|
||||
json!({
|
||||
"contents": [{ "parts": [{ "text": prompt }] }]
|
||||
})
|
||||
}
|
||||
|
||||
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()
|
||||
);
|
||||
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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
215
src/gemini/errors.rs
Normal file
215
src/gemini/errors.rs
Normal file
@@ -0,0 +1,215 @@
|
||||
use reqwest::Response;
|
||||
use serde::Deserialize;
|
||||
use std::time::Duration;
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum GeminiError {
|
||||
#[error("API rate limit exceeded. Retry after {retry_after} seconds")]
|
||||
RateLimitExceeded { retry_after: u32 },
|
||||
|
||||
#[error("Quota exceeded. Usage limit reached: {limit}")]
|
||||
QuotaExceeded { limit: String },
|
||||
|
||||
#[error("Model '{model}' not found or unavailable")]
|
||||
ModelNotFound { model: String },
|
||||
|
||||
#[error("Invalid API key. Please check your GEMINI_API_KEY")]
|
||||
InvalidApiKey,
|
||||
|
||||
#[error("Content policy violation: {reason}")]
|
||||
ContentPolicyViolation { reason: String },
|
||||
|
||||
#[error("Invalid request: {details}")]
|
||||
InvalidRequest { details: String },
|
||||
|
||||
#[error("Network error: {0}")]
|
||||
NetworkError(#[from] reqwest::Error),
|
||||
|
||||
#[error("Invalid response format: {0}")]
|
||||
InvalidResponse(String),
|
||||
|
||||
#[error("API error (HTTP {status}): {message}")]
|
||||
ApiError { status: u16, message: String },
|
||||
|
||||
#[error("Service temporarily unavailable: {reason}")]
|
||||
ServiceUnavailable { reason: String },
|
||||
|
||||
#[error("Request timeout after {seconds} seconds")]
|
||||
Timeout { seconds: u64 },
|
||||
|
||||
#[error("JSON serialization/deserialization error: {0}")]
|
||||
SerializationError(#[from] serde_json::Error),
|
||||
|
||||
#[error("Internal server error: {details}")]
|
||||
InternalError { details: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct GeminiErrorResponse {
|
||||
error: GeminiErrorDetail,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct GeminiErrorDetail {
|
||||
#[allow(dead_code)]
|
||||
code: i32,
|
||||
message: String,
|
||||
status: String,
|
||||
#[serde(default)]
|
||||
details: Vec<GeminiErrorDetailInfo>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct GeminiErrorDetailInfo {
|
||||
#[serde(rename = "@type")]
|
||||
#[allow(dead_code)]
|
||||
error_type: String,
|
||||
#[serde(rename = "retryDelay")]
|
||||
retry_delay: Option<String>,
|
||||
quota_limit: Option<String>,
|
||||
#[allow(dead_code)]
|
||||
quota_metro: Option<String>,
|
||||
}
|
||||
|
||||
impl GeminiError {
|
||||
/// Parse HTTP response and convert to appropriate GeminiError
|
||||
pub async fn from_response(response: Response) -> Self {
|
||||
let status = response.status();
|
||||
|
||||
let error_text = match response.text().await {
|
||||
Ok(text) => text,
|
||||
Err(e) => {
|
||||
return GeminiError::NetworkError(e);
|
||||
}
|
||||
};
|
||||
|
||||
if let Ok(gemini_error) = serde_json::from_str::<GeminiErrorResponse>(&error_text) {
|
||||
return Self::from_gemini_error(gemini_error.error, status.as_u16());
|
||||
}
|
||||
|
||||
Self::from_status_code(status, &error_text)
|
||||
}
|
||||
|
||||
fn from_gemini_error(error_detail: GeminiErrorDetail, status: u16) -> Self {
|
||||
let details = error_detail.details;
|
||||
|
||||
match error_detail.status.as_str() {
|
||||
"RESOURCE_EXHAUSTED" => {
|
||||
if let Some(retry_info) = details.iter().find(|d| d.retry_delay.is_some())
|
||||
&& let Some(retry_delay) = &retry_info.retry_delay
|
||||
&& let Ok(seconds) = retry_delay.parse::<u32>()
|
||||
{
|
||||
return GeminiError::RateLimitExceeded {
|
||||
retry_after: seconds,
|
||||
};
|
||||
}
|
||||
|
||||
if let Some(quota_info) = details.iter().find(|d| d.quota_limit.is_some()) {
|
||||
let limit = quota_info.quota_limit.as_deref().unwrap_or("unknown");
|
||||
return GeminiError::QuotaExceeded {
|
||||
limit: limit.to_string(),
|
||||
};
|
||||
}
|
||||
|
||||
GeminiError::QuotaExceeded {
|
||||
limit: "usage limit".to_string(),
|
||||
}
|
||||
}
|
||||
"NOT_FOUND" => {
|
||||
// Extract model name from message if possible
|
||||
let model = extract_model_name(&error_detail.message);
|
||||
GeminiError::ModelNotFound { model }
|
||||
}
|
||||
"UNAUTHENTICATED" => GeminiError::InvalidApiKey,
|
||||
"PERMISSION_DENIED" => {
|
||||
if error_detail.message.to_lowercase().contains("policy") {
|
||||
GeminiError::ContentPolicyViolation {
|
||||
reason: error_detail.message,
|
||||
}
|
||||
} else {
|
||||
GeminiError::InvalidRequest {
|
||||
details: error_detail.message,
|
||||
}
|
||||
}
|
||||
}
|
||||
"INVALID_ARGUMENT" => GeminiError::InvalidRequest {
|
||||
details: error_detail.message,
|
||||
},
|
||||
"UNAVAILABLE" => GeminiError::ServiceUnavailable {
|
||||
reason: error_detail.message,
|
||||
},
|
||||
"DEADLINE_EXCEEDED" => GeminiError::Timeout { seconds: 60 },
|
||||
"INTERNAL" => GeminiError::InternalError {
|
||||
details: error_detail.message,
|
||||
},
|
||||
_ => GeminiError::ApiError {
|
||||
status,
|
||||
message: error_detail.message,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn from_status_code(status: reqwest::StatusCode, error_text: &str) -> Self {
|
||||
match status.as_u16() {
|
||||
400 => GeminiError::InvalidRequest {
|
||||
details: error_text.to_string(),
|
||||
},
|
||||
401 => GeminiError::InvalidApiKey,
|
||||
403 => GeminiError::ContentPolicyViolation {
|
||||
reason: error_text.to_string(),
|
||||
},
|
||||
404 => GeminiError::ModelNotFound {
|
||||
model: "unknown".to_string(),
|
||||
},
|
||||
429 => GeminiError::RateLimitExceeded { retry_after: 60 },
|
||||
500 => GeminiError::InternalError {
|
||||
details: error_text.to_string(),
|
||||
},
|
||||
502..=504 => GeminiError::ServiceUnavailable {
|
||||
reason: error_text.to_string(),
|
||||
},
|
||||
_ => GeminiError::ApiError {
|
||||
status: status.as_u16(),
|
||||
message: error_text.to_string(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if this error is retryable
|
||||
pub fn is_retryable(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
GeminiError::RateLimitExceeded { .. }
|
||||
| GeminiError::ServiceUnavailable { .. }
|
||||
| GeminiError::Timeout { .. }
|
||||
| GeminiError::NetworkError(_)
|
||||
| GeminiError::InternalError { .. }
|
||||
)
|
||||
}
|
||||
|
||||
/// Get retry delay for retryable errors
|
||||
pub fn retry_delay(&self) -> Option<Duration> {
|
||||
match self {
|
||||
GeminiError::RateLimitExceeded { retry_after } => {
|
||||
Some(Duration::from_secs(*retry_after as u64))
|
||||
}
|
||||
GeminiError::ServiceUnavailable { .. } => Some(Duration::from_secs(30)),
|
||||
GeminiError::NetworkError(_) => Some(Duration::from_secs(5)),
|
||||
GeminiError::Timeout { .. } => Some(Duration::from_secs(10)),
|
||||
GeminiError::InternalError { .. } => Some(Duration::from_secs(15)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_model_name(message: &str) -> String {
|
||||
// Try to extract model name from error message
|
||||
// Example: "Model 'gemini-1.5-flash' not found"
|
||||
if let Some(start) = message.find('\'')
|
||||
&& let Some(end) = message[start + 1..].find('\'')
|
||||
{
|
||||
return message[start + 1..start + 1 + end].to_string();
|
||||
}
|
||||
"unknown".to_string()
|
||||
}
|
||||
8
src/gemini/mod.rs
Normal file
8
src/gemini/mod.rs
Normal file
@@ -0,0 +1,8 @@
|
||||
pub mod client;
|
||||
pub mod errors;
|
||||
pub mod prompt;
|
||||
pub mod types;
|
||||
|
||||
pub use client::GeminiClient;
|
||||
pub use errors::GeminiError;
|
||||
pub use types::{Candidate, Content, FileCategoryResponse, GeminiResponse, OrganizationPlanResponse, Part};
|
||||
51
src/gemini/prompt.rs
Normal file
51
src/gemini/prompt.rs
Normal file
@@ -0,0 +1,51 @@
|
||||
use crate::models::{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>,
|
||||
}
|
||||
Reference in New Issue
Block a user