use anyhow::{Context, Result}; use chrono::Utc; use clap::Parser; use image_api::ai::ollama::OllamaClient; use image_api::database::calendar_dao::{InsertCalendarEvent, SqliteCalendarEventDao}; use image_api::parsers::ical_parser::parse_ics_file; use log::{error, info}; use std::sync::{Arc, Mutex}; // Import the trait to use its methods use image_api::database::CalendarEventDao; #[derive(Parser, Debug)] #[command(author, version, about = "Import Google Takeout Calendar data", long_about = None)] struct Args { /// Path to the .ics calendar file #[arg(short, long)] path: String, /// Generate embeddings for calendar events (slower but enables semantic search) #[arg(long, default_value = "false")] generate_embeddings: bool, /// Skip events that already exist in the database #[arg(long, default_value = "true")] skip_existing: bool, /// Batch size for embedding generation #[arg(long, default_value = "128")] batch_size: usize, } #[tokio::main] async fn main() -> Result<()> { dotenv::dotenv().ok(); env_logger::init(); let args = Args::parse(); info!("Parsing calendar file: {}", args.path); let events = parse_ics_file(&args.path).context("Failed to parse .ics file")?; info!("Found {} calendar events", events.len()); let context = opentelemetry::Context::current(); let ollama = if args.generate_embeddings { let primary_url = dotenv::var("OLLAMA_PRIMARY_URL") .or_else(|_| dotenv::var("OLLAMA_URL")) .unwrap_or_else(|_| "http://localhost:11434".to_string()); let fallback_url = dotenv::var("OLLAMA_FALLBACK_URL").ok(); let primary_model = dotenv::var("OLLAMA_PRIMARY_MODEL") .or_else(|_| dotenv::var("OLLAMA_MODEL")) .unwrap_or_else(|_| "nomic-embed-text:v1.5".to_string()); let fallback_model = dotenv::var("OLLAMA_FALLBACK_MODEL").ok(); Some(OllamaClient::new( primary_url, fallback_url, primary_model, fallback_model, )) } else { None }; let inserted_count = Arc::new(Mutex::new(0)); let skipped_count = Arc::new(Mutex::new(0)); let error_count = Arc::new(Mutex::new(0)); // Process events in batches // Can't use rayon with async, so process sequentially for event in &events { let mut dao_instance = SqliteCalendarEventDao::new(); // Check if event exists if args.skip_existing && let Ok(exists) = dao_instance.event_exists( &context, event.event_uid.as_deref().unwrap_or(""), event.start_time, ) && exists { *skipped_count.lock().unwrap() += 1; continue; } // Generate embedding if requested (blocking call) let embedding = if let Some(ref ollama_client) = ollama { let text = format!( "{} {} {}", event.summary, event.description.as_deref().unwrap_or(""), event.location.as_deref().unwrap_or("") ); match tokio::task::block_in_place(|| { tokio::runtime::Handle::current() .block_on(async { ollama_client.generate_embedding(&text).await }) }) { Ok(emb) => Some(emb), Err(e) => { error!( "Failed to generate embedding for event '{}': {}", event.summary, e ); None } } } else { None }; // Insert into database let insert_event = InsertCalendarEvent { event_uid: event.event_uid.clone(), summary: event.summary.clone(), description: event.description.clone(), location: event.location.clone(), start_time: event.start_time, end_time: event.end_time, all_day: event.all_day, organizer: event.organizer.clone(), attendees: if event.attendees.is_empty() { None } else { Some(serde_json::to_string(&event.attendees).unwrap_or_default()) }, embedding, created_at: Utc::now().timestamp(), source_file: Some(args.path.clone()), }; match dao_instance.store_event(&context, insert_event) { Ok(_) => { *inserted_count.lock().unwrap() += 1; if *inserted_count.lock().unwrap() % 100 == 0 { info!("Imported {} events...", *inserted_count.lock().unwrap()); } } Err(e) => { error!("Failed to store event '{}': {:?}", event.summary, e); *error_count.lock().unwrap() += 1; } } } let final_inserted = *inserted_count.lock().unwrap(); let final_skipped = *skipped_count.lock().unwrap(); let final_errors = *error_count.lock().unwrap(); info!("\n=== Import Summary ==="); info!("Total events found: {}", events.len()); info!("Successfully inserted: {}", final_inserted); info!("Skipped (already exist): {}", final_skipped); info!("Errors: {}", final_errors); if args.generate_embeddings { info!("Embeddings were generated for semantic search"); } else { info!("No embeddings generated (use --generate-embeddings to enable semantic search)"); } Ok(()) }