Files
ImageApi/src/bin/import_calendar.rs
Cameron af35a996a3 Cleanup unused message embedding code
Fixup some warnings
2026-01-14 13:33:36 -05:00

167 lines
5.5 KiB
Rust

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