From f3e34dc2590a7cbadc4cd073f1ee8a63c386b17c Mon Sep 17 00:00:00 2001 From: scald <1215913+scald@users.noreply.github.com> Date: Mon, 25 Nov 2024 14:22:38 +0200 Subject: [PATCH] formatting --- src/errors.rs | 14 +++--- src/lib.rs | 16 +++---- src/main.rs | 35 ++++++++------- src/processor.rs | 98 +++++++++++++++++++++++------------------ src/prompts.rs | 79 ++++++++++++++++++--------------- src/providers/mod.rs | 10 ++--- src/providers/ollama.rs | 18 +++++--- src/providers/openai.rs | 36 +++++++++------ src/utils.rs | 27 +++++++----- 9 files changed, 184 insertions(+), 149 deletions(-) diff --git a/src/errors.rs b/src/errors.rs index e78b6ee..7a05ceb 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -4,22 +4,22 @@ use thiserror::Error; pub enum ProcessorError { #[error("Failed to load image: {0}")] ImageLoadError(#[from] image::error::ImageError), - + #[error("AI Provider error: {0}")] AIProviderError(String), - + #[error("Failed to encode/decode base64: {0}")] Base64Error(String), - + #[error("Environment variable error: {0}")] EnvError(#[from] std::env::VarError), - + #[error("Network request failed: {0}")] RequestError(#[from] reqwest::Error), - + #[error("Invalid API response: {0}")] ResponseParseError(String), - + #[error("Thumbnail generation failed: {0}")] ThumbnailError(String), -} \ No newline at end of file +} diff --git a/src/lib.rs b/src/lib.rs index 44d9af7..c01a8a9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,16 +1,16 @@ //! eyeris is a high-performance image analysis service. -//! +//! //! # Features -//! +//! //! - Multiple AI provider support (OpenAI, Ollama) //! - Image optimization and processing //! - Customizable analysis formats -//! +//! //! # Example -//! +//! //! ```rust,no_run //! use eyeris::{ImageProcessor, AIProvider}; -//! +//! //! #[tokio::main] //! async fn main() { //! let processor = ImageProcessor::new( @@ -26,14 +26,14 @@ //! } //! ``` +pub mod errors; pub mod processor; pub mod prompts; pub mod providers; -pub mod errors; pub mod utils; // Re-export commonly used types +pub use errors::ProcessorError; pub use processor::ImageProcessor; pub use prompts::{ImagePrompt, PromptFormat}; -pub use errors::ProcessorError; -pub use providers::AIProvider; \ No newline at end of file +pub use providers::AIProvider; diff --git a/src/main.rs b/src/main.rs index 97ebbf9..fe39ea7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,4 @@ +use anyhow::Result; use axum::{ extract::{Multipart, Query, State}, http::StatusCode, @@ -5,18 +6,13 @@ use axum::{ routing::post, Router, }; +use bytes::Bytes; +use eyeris::{processor::ImageProcessor, prompts::PromptFormat, providers::AIProvider}; use serde::{Deserialize, Serialize}; use std::sync::Arc; -use tokio::sync::Semaphore; -use anyhow::Result; -use bytes::Bytes; use std::sync::OnceLock; use std::time::Instant; -use eyeris::{ - providers::AIProvider, - processor::ImageProcessor, - prompts::PromptFormat, -}; +use tokio::sync::Semaphore; // Create a static processor pool static PROCESSOR_POOL: OnceLock> = OnceLock::new(); @@ -82,7 +78,7 @@ async fn main() -> Result<()> { // Start server let addr = "0.0.0.0:3000"; tracing::info!("Starting server on {}", addr); - + let listener = tokio::net::TcpListener::bind(addr).await?; axum::serve(listener, app).await?; @@ -115,7 +111,10 @@ async fn process_image( }); let permit_start = Instant::now(); - let _permit = state.semaphore.acquire().await + let _permit = state + .semaphore + .acquire() + .await .map_err(|_| StatusCode::SERVICE_UNAVAILABLE)?; tracing::info!( duration_ms = permit_start.elapsed().as_millis(), @@ -123,8 +122,10 @@ async fn process_image( ); let multipart_start = Instant::now(); - let image_data: Bytes = if let Some(field) = multipart.next_field().await - .map_err(|_| StatusCode::BAD_REQUEST)? + let image_data: Bytes = if let Some(field) = multipart + .next_field() + .await + .map_err(|_| StatusCode::BAD_REQUEST)? { if field.name().unwrap_or("") != "image" { return Ok(Json(ProcessResponse { @@ -163,24 +164,24 @@ async fn process_image( Ok(analysis) => { // Get current token stats let token_stats = processor.token_stats.read(); - + Ok(Json(ProcessResponse { success: true, message: "Image processed successfully".to_string(), - data: Some(ProcessedData { + data: Some(ProcessedData { analysis, token_usage: TokenUsage { prompt_tokens: token_stats.prompt_tokens, completion_tokens: token_stats.completion_tokens, total_tokens: token_stats.total_tokens, - } + }, }), })) - }, + } Err(e) => Ok(Json(ProcessResponse { success: false, message: format!("Processing failed: {}", e), data: None, })), } -} \ No newline at end of file +} diff --git a/src/processor.rs b/src/processor.rs index 19d0c2b..ff2382c 100644 --- a/src/processor.rs +++ b/src/processor.rs @@ -1,12 +1,12 @@ -use image::{ImageBuffer, DynamicImage}; -use rayon::prelude::*; +use crate::errors::ProcessorError; +use crate::prompts::{ImagePrompt, PromptFormat}; +use crate::providers::{AIProvider, OllamaProvider, OpenAIProvider, Provider}; use base64::Engine; use image::codecs::jpeg::JpegEncoder; use image::imageops::FilterType; -use crate::prompts::{ImagePrompt, PromptFormat}; +use image::{DynamicImage, ImageBuffer}; +use rayon::prelude::*; use std::time::Instant; -use crate::providers::{Provider, OpenAIProvider, OllamaProvider, AIProvider}; -use crate::errors::ProcessorError; #[derive(Debug, Clone, Default)] pub struct TokenStats { @@ -44,7 +44,7 @@ impl ImageProcessor { /// ``` pub fn new(provider: AIProvider, model: Option, format: Option) -> Self { let token_stats = std::sync::Arc::new(parking_lot::RwLock::new(TokenStats::default())); - + let provider: Box = match provider { AIProvider::OpenAI => Box::new(OpenAIProvider::new(model, token_stats.clone())), AIProvider::Ollama => Box::new(OllamaProvider::new(model)), @@ -59,7 +59,7 @@ impl ImageProcessor { pub async fn process(&self, image_data: &[u8]) -> Result { let start = Instant::now(); - + // Pre-allocate the base64 string let base64_start = Instant::now(); let base64_image = base64::engine::general_purpose::STANDARD.encode(image_data); @@ -78,7 +78,7 @@ impl ImageProcessor { }; let (analysis, _thumbnail_result) = tokio::join!(analysis_future, thumbnail_future); - + tracing::info!( duration_ms = parallel_start.elapsed().as_millis(), "Parallel processing completed" @@ -94,20 +94,21 @@ impl ImageProcessor { async fn analyze_image(&self, base64_image: &str) -> Result { let start = Instant::now(); - + // Get the length before moving the string let original_size = base64_image.len(); - + // Clone the string before moving into spawn_blocking let base64_image_owned = base64_image.to_string(); let optimized_image = tokio::task::spawn_blocking(move || { // Use base64_image_owned instead of base64_image - let image_data = base64::engine::general_purpose::STANDARD.decode(&base64_image_owned) + let image_data = base64::engine::general_purpose::STANDARD + .decode(&base64_image_owned) .map_err(|e| ProcessorError::Base64Error(e.to_string()))?; - + // Load image let img = image::load_from_memory(&image_data)?; - + // Resize to smaller dimensions while maintaining aspect ratio // Most vision models don't need images larger than 768px let max_dim = 768; @@ -122,22 +123,27 @@ impl ImageProcessor { // Convert to RGB and compress as JPEG with moderate quality let mut buffer = Vec::new(); - let mut encoder = JpegEncoder::new_with_quality(&mut buffer, 10); - encoder.encode( - resized.to_rgb8().as_raw(), - resized.width(), - resized.height(), - image::ColorType::Rgb8 - ).map_err(|e| ProcessorError::ImageLoadError(e))?; + let mut encoder = JpegEncoder::new_with_quality(&mut buffer, 10); + encoder + .encode( + resized.to_rgb8().as_raw(), + resized.width(), + resized.height(), + image::ColorType::Rgb8, + ) + .map_err(|e| ProcessorError::ImageLoadError(e))?; // Convert back to base64 Ok::<_, ProcessorError>(base64::engine::general_purpose::STANDARD.encode(&buffer)) - }).await.map_err(|e| ProcessorError::ThumbnailError(e.to_string()))??; + }) + .await + .map_err(|e| ProcessorError::ThumbnailError(e.to_string()))??; tracing::info!( original_size = original_size, optimized_size = optimized_image.len(), - reduction_percent = ((original_size - optimized_image.len()) as f32 / original_size as f32 * 100.0), + reduction_percent = + ((original_size - optimized_image.len()) as f32 / original_size as f32 * 100.0), "Image optimization completed" ); @@ -156,7 +162,7 @@ impl ImageProcessor { async fn create_thumbnail(&self, image: DynamicImage) -> Result, ProcessorError> { let start = Instant::now(); - + let result = tokio::task::spawn_blocking(move || { let resize_start = Instant::now(); let thumbnail = image.resize(300, 300, FilterType::Triangle); @@ -165,23 +171,26 @@ impl ImageProcessor { duration_ms = resize_start.elapsed().as_millis(), "Image resize completed" ); - + let enhance_start = Instant::now(); let width = rgb_image.width() as usize; let height = rgb_image.height() as usize; - - let enhanced_pixels: Vec<_> = rgb_image.chunks_exact(width * 3) + + let enhanced_pixels: Vec<_> = rgb_image + .chunks_exact(width * 3) .enumerate() .par_bridge() .flat_map(|(y, row)| { - (0..width).map(move |x| { - let pixel = image::Rgb([ - (row[x * 3] as f32 * 1.1).min(255.0) as u8, - (row[x * 3 + 1] as f32 * 1.1).min(255.0) as u8, - (row[x * 3 + 2] as f32 * 1.1).min(255.0) as u8, - ]); - (x as u32, y as u32, pixel) - }).collect::>() + (0..width) + .map(move |x| { + let pixel = image::Rgb([ + (row[x * 3] as f32 * 1.1).min(255.0) as u8, + (row[x * 3 + 1] as f32 * 1.1).min(255.0) as u8, + (row[x * 3 + 2] as f32 * 1.1).min(255.0) as u8, + ]); + (x as u32, y as u32, pixel) + }) + .collect::>() }) .collect(); @@ -198,20 +207,23 @@ impl ImageProcessor { let mut output = Vec::with_capacity(width * height * 3); let mut encoder = JpegEncoder::new_with_quality(&mut output, 85); - encoder.encode( - enhanced.as_raw(), - enhanced.width(), - enhanced.height(), - image::ColorType::Rgb8 - ).map_err(|e| ProcessorError::ThumbnailError(e.to_string()))?; + encoder + .encode( + enhanced.as_raw(), + enhanced.width(), + enhanced.height(), + image::ColorType::Rgb8, + ) + .map_err(|e| ProcessorError::ThumbnailError(e.to_string()))?; tracing::info!( duration_ms = buffer_start.elapsed().as_millis(), "Buffer creation and JPEG encoding completed" ); - + Ok(output) - }).await + }) + .await .map_err(|e| ProcessorError::ThumbnailError(e.to_string()))?; tracing::info!( @@ -221,4 +233,4 @@ impl ImageProcessor { result } -} \ No newline at end of file +} diff --git a/src/prompts.rs b/src/prompts.rs index 8203e44..dd40e81 100644 --- a/src/prompts.rs +++ b/src/prompts.rs @@ -20,14 +20,16 @@ pub enum PromptFormat { #[serde(rename_all = "snake_case")] pub enum ContentCategory { // Digital Content - Screenshot { platform: Option }, + Screenshot { + platform: Option, + }, UserInterface, SocialMediaPost, DigitalArt, Website, Software, VideoGame, - + // Documents Document, Receipt, @@ -36,7 +38,7 @@ pub enum ContentCategory { Form, Identification, Certificate, - + // Visual Content Photo, Artwork, @@ -45,7 +47,7 @@ pub enum ContentCategory { Comic, Advertisement, Poster, - + // Instructional Recipe, Tutorial, @@ -54,7 +56,7 @@ pub enum ContentCategory { Schematic, Manual, Guide, - + // Data Visualization Chart, Graph, @@ -63,14 +65,14 @@ pub enum ContentCategory { Timeline, Flowchart, MindMap, - + // Location/Space Map, FloorPlan, Architecture, Landscape, Satellite, - + // Special Purpose Medical, Scientific, @@ -78,7 +80,7 @@ pub enum ContentCategory { Educational, Legal, Financial, - + // Dynamic - for discovered categories Discovered { name: String, @@ -162,17 +164,17 @@ impl ImagePrompt { format, config, }; - + // Add dynamic discovery instructions let dynamic_text = prompt.add_dynamic_discovery_prompt(); prompt.text.push_str("\n\n"); prompt.text.push_str(&dynamic_text); - + prompt } fn get_concise_prompt(config: &AnalysisConfig) -> String { let mut prompt = "Analyze this image and describe its contents concisely.".to_string(); - + if config.extract_text { prompt.push_str(" Extract any visible text."); } @@ -399,17 +401,20 @@ impl ImagePrompt { - Note accessibility considerations - Identify platform-specific patterns - Map navigation structure -- Flag any security/privacy concerns"#.to_string(); +- Flag any security/privacy concerns"# + .to_string(); if let Some(platform_name) = platform { - instructions.push_str(&format!("\n\nSpecific to {platform_name}: + instructions.push_str(&format!( + "\n\nSpecific to {platform_name}: - Check platform-specific design guidelines - Identify standard platform components -- Note any platform-specific conventions")); +- Note any platform-specific conventions" + )); } instructions - }, + } // ... [rest of the category matches from your original code] _ => r#"Perform comprehensive analysis considering: - Primary purpose @@ -420,7 +425,8 @@ impl ImagePrompt { - Practical applications - Quality indicators - Notable patterns -- Unique characteristics"#.to_string(), +- Unique characteristics"# + .to_string(), } } @@ -492,7 +498,7 @@ mod tests { for format in formats { let prompt = ImagePrompt::new(format.clone()); assert!(!prompt.text.is_empty()); - + // Test with specific configuration let config = AnalysisConfig { extract_text: true, @@ -509,17 +515,19 @@ mod tests { cultural_analysis: true, technical_details: true, accessibility_analysis: true, - content_category: Some(ContentCategory::Screenshot { platform: Some("iOS".to_string()) }), + content_category: Some(ContentCategory::Screenshot { + platform: Some("iOS".to_string()), + }), custom_traits: vec![], }; - + let prompt_with_config = ImagePrompt::with_config(format, config); assert!(!prompt_with_config.text.is_empty()); - + // Test OpenAI content generation let openai_content = prompt.to_openai_content(); assert!(openai_content.is_array()); - + // Test Ollama prompt generation let ollama_prompt = prompt.to_ollama_prompt(); assert!(!ollama_prompt.is_empty()); @@ -529,7 +537,9 @@ mod tests { #[test] fn test_category_specific_prompts() { let categories = vec![ - ContentCategory::Screenshot { platform: Some("iOS".to_string()) }, + ContentCategory::Screenshot { + platform: Some("iOS".to_string()), + }, ContentCategory::Recipe, ContentCategory::Document, ContentCategory::Map, @@ -540,11 +550,11 @@ mod tests { content_category: Some(category), ..Default::default() }; - + let prompt = ImagePrompt::with_config(PromptFormat::Json, config); assert!(prompt.text.contains("For this")); assert!(!prompt.text.is_empty()); - + // Verify dynamic discovery prompt is included assert!(prompt.text.contains("Pattern Recognition")); assert!(prompt.text.contains("Innovation Detection")); @@ -559,9 +569,9 @@ mod tests { cultural_analysis: true, ..Default::default() }; - + let prompt = ImagePrompt::with_config(PromptFormat::Json, config); - + // Check for dynamic analysis elements assert!(prompt.text.contains("dynamic_extensions")); assert!(prompt.text.contains("discovered_categories")); @@ -571,13 +581,10 @@ mod tests { #[test] fn test_custom_traits() { let config = AnalysisConfig { - custom_traits: vec![ - "brand_safety".to_string(), - "viral_potential".to_string(), - ], + custom_traits: vec!["brand_safety".to_string(), "viral_potential".to_string()], ..Default::default() }; - + let prompt = ImagePrompt::with_config(PromptFormat::Detailed, config); assert!(prompt.text.contains("brand_safety") || prompt.text.contains("viral_potential")); } @@ -586,11 +593,11 @@ mod tests { fn test_platform_specific_screenshot() { let config = AnalysisConfig { content_category: Some(ContentCategory::Screenshot { - platform: Some("iOS".to_string()) + platform: Some("iOS".to_string()), }), ..Default::default() }; - + let prompt = ImagePrompt::with_config(PromptFormat::Detailed, config); assert!(prompt.text.contains("iOS")); assert!(prompt.text.contains("platform-specific")); @@ -601,14 +608,14 @@ mod tests { let prompt = ImagePrompt::new(PromptFormat::Json); let serialized = serde_json::to_string(&prompt).unwrap(); assert!(!serialized.is_empty()); - + // Test OpenAI format let openai_content = prompt.to_openai_content(); assert!(openai_content.is_array()); assert_eq!(openai_content.as_array().unwrap().len(), 2); - + // Test Ollama format let ollama_prompt = prompt.to_ollama_prompt(); assert!(!ollama_prompt.is_empty()); } -} \ No newline at end of file +} diff --git a/src/providers/mod.rs b/src/providers/mod.rs index e9d2ac6..dc8ef22 100644 --- a/src/providers/mod.rs +++ b/src/providers/mod.rs @@ -1,10 +1,10 @@ -mod openai; mod ollama; +mod openai; -pub use openai::OpenAIProvider; -pub use ollama::OllamaProvider; -use async_trait::async_trait; use crate::errors::ProcessorError; +use async_trait::async_trait; +pub use ollama::OllamaProvider; +pub use openai::OpenAIProvider; #[derive(Debug, Clone, Copy)] pub enum AIProvider { @@ -15,4 +15,4 @@ pub enum AIProvider { #[async_trait] pub trait Provider: Send + Sync { async fn analyze(&self, base64_image: &str, prompt: &str) -> Result; -} \ No newline at end of file +} diff --git a/src/providers/ollama.rs b/src/providers/ollama.rs index 0aa48c9..00f4b5e 100644 --- a/src/providers/ollama.rs +++ b/src/providers/ollama.rs @@ -39,7 +39,8 @@ impl Provider for OllamaProvider { images: vec![base64_image.to_string()], }; - let response = self.client + let response = self + .client .post("http://localhost:11434/api/generate") .json(&ollama_request) .send() @@ -47,18 +48,19 @@ impl Provider for OllamaProvider { let status = response.status(); if !status.is_success() { - let error_text = response.text().await + let error_text = response + .text() + .await .unwrap_or_else(|_| "Failed to get error message".to_string()); return Err(ProcessorError::AIProviderError(format!( "Ollama API request failed with status {}: {}", - status, - error_text + status, error_text ))); } // Ollama returns streaming responses, so we need to collect all response chunks let text = response.text().await?; - + // Parse each line as a separate JSON response let mut full_response = String::new(); for line in text.lines() { @@ -68,9 +70,11 @@ impl Provider for OllamaProvider { } if full_response.is_empty() { - return Err(ProcessorError::ResponseParseError("Empty response from Ollama".to_string())); + return Err(ProcessorError::ResponseParseError( + "Empty response from Ollama".to_string(), + )); } Ok(full_response) } -} \ No newline at end of file +} diff --git a/src/providers/openai.rs b/src/providers/openai.rs index 49533c2..79cfe5e 100644 --- a/src/providers/openai.rs +++ b/src/providers/openai.rs @@ -1,12 +1,12 @@ use super::Provider; use crate::errors::ProcessorError; +use crate::processor::TokenStats; use async_trait::async_trait; +use parking_lot::RwLock; use reqwest::Client; use serde::Deserialize; use serde_json::json; use std::sync::Arc; -use parking_lot::RwLock; -use crate::processor::TokenStats; #[derive(Debug, Deserialize)] struct OpenAIResponse { @@ -50,8 +50,7 @@ impl OpenAIProvider { #[async_trait] impl Provider for OpenAIProvider { async fn analyze(&self, base64_image: &str, prompt: &str) -> Result { - let api_key = std::env::var("OPENAI_API_KEY") - .map_err(|e| ProcessorError::EnvError(e))?; + let api_key = std::env::var("OPENAI_API_KEY").map_err(|e| ProcessorError::EnvError(e))?; let request_body = json!({ "model": self.model, @@ -73,7 +72,8 @@ impl Provider for OpenAIProvider { "max_tokens": 1000 }); - let response = self.client + let response = self + .client .post("https://api.openai.com/v1/chat/completions") .header("Authorization", format!("Bearer {}", api_key)) .json(&request_body) @@ -82,18 +82,23 @@ impl Provider for OpenAIProvider { let status = response.status(); if !status.is_success() { - let error_text = response.text().await + let error_text = response + .text() + .await .unwrap_or_else(|_| "Failed to get error message".to_string()); return Err(ProcessorError::AIProviderError(format!( "OpenAI API request failed with status {}: {}", - status, - error_text + status, error_text ))); } let response_text = response.text().await?; - let response: OpenAIResponse = serde_json::from_str(&response_text) - .map_err(|e| ProcessorError::ResponseParseError(format!("Failed to parse OpenAI response: {}. Response text: {}", e, response_text)))?; + let response: OpenAIResponse = serde_json::from_str(&response_text).map_err(|e| { + ProcessorError::ResponseParseError(format!( + "Failed to parse OpenAI response: {}. Response text: {}", + e, response_text + )) + })?; // Update token stats if usage information is available if let Some(usage) = response.usage { @@ -101,7 +106,7 @@ impl Provider for OpenAIProvider { stats.prompt_tokens += usage.prompt_tokens; stats.completion_tokens += usage.completion_tokens; stats.total_tokens += usage.total_tokens; - + tracing::info!( prompt_tokens = usage.prompt_tokens, completion_tokens = usage.completion_tokens, @@ -110,13 +115,16 @@ impl Provider for OpenAIProvider { ); } - let analysis = response.choices + let analysis = response + .choices .first() - .ok_or_else(|| ProcessorError::ResponseParseError("No choices in response".to_string()))? + .ok_or_else(|| { + ProcessorError::ResponseParseError("No choices in response".to_string()) + })? .message .content .clone(); Ok(analysis) } -} \ No newline at end of file +} diff --git a/src/utils.rs b/src/utils.rs index 9f526aa..b1f313f 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,24 +1,27 @@ +use crate::errors::ProcessorError; use image::{DynamicImage, ImageBuffer, Rgb}; use rayon::prelude::*; -use crate::errors::ProcessorError; pub fn enhance_image(img: &DynamicImage) -> Result, Vec>, ProcessorError> { let rgb_image = img.to_rgb8(); let width = rgb_image.width() as usize; let height = rgb_image.height() as usize; - - let enhanced_pixels: Vec<_> = rgb_image.chunks_exact(width * 3) + + let enhanced_pixels: Vec<_> = rgb_image + .chunks_exact(width * 3) .enumerate() .par_bridge() .flat_map(|(y, row)| { - (0..width).map(move |x| { - let pixel = Rgb([ - (row[x * 3] as f32 * 1.1).min(255.0) as u8, - (row[x * 3 + 1] as f32 * 1.1).min(255.0) as u8, - (row[x * 3 + 2] as f32 * 1.1).min(255.0) as u8, - ]); - (x as u32, y as u32, pixel) - }).collect::>() + (0..width) + .map(move |x| { + let pixel = Rgb([ + (row[x * 3] as f32 * 1.1).min(255.0) as u8, + (row[x * 3 + 1] as f32 * 1.1).min(255.0) as u8, + (row[x * 3 + 2] as f32 * 1.1).min(255.0) as u8, + ]); + (x as u32, y as u32, pixel) + }) + .collect::>() }) .collect(); @@ -28,4 +31,4 @@ pub fn enhance_image(img: &DynamicImage) -> Result, Vec> } Ok(enhanced) -} \ No newline at end of file +}