diff --git a/src/ai/claude.rs b/src/ai/claude.rs index 42582d6..69dfb80 100644 --- a/src/ai/claude.rs +++ b/src/ai/claude.rs @@ -4,7 +4,9 @@ use serde::{Deserialize, Serialize}; use crate::model::{ContentDescription, FileSummary, ProposedGroup}; -use super::{AiProvider, DescribeContext}; +use super::{ + AiProvider, DescribeContext, DescribePayload, DescribeRequest, +}; const LOCAL_API_KEY: &str = "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"; @@ -16,8 +18,12 @@ pub struct ClaudeProvider { model: String, base_url: String, max_retries: usize, + poll_interval: std::time::Duration, } +const DEFAULT_BATCH_POLL_INTERVAL: std::time::Duration = + std::time::Duration::from_secs(5); + #[derive(Serialize)] struct Message { role: &'static str, @@ -107,6 +113,42 @@ struct ApiResponse { stop_reason: Option, } +#[derive(Serialize)] +struct BatchRequestItem { + custom_id: String, + params: ApiRequest, +} + +#[derive(Serialize)] +struct BatchSubmitBody { + requests: Vec, +} + +#[derive(Deserialize)] +struct BatchStatus { + id: String, + processing_status: String, +} + +#[derive(Deserialize)] +struct BatchResultLine { + custom_id: String, + result: BatchResult, +} + +#[derive(Deserialize)] +#[serde(tag = "type")] +enum BatchResult { + #[serde(rename = "succeeded")] + Succeeded { message: ApiResponse }, + #[serde(rename = "errored")] + Errored { error: serde_json::Value }, + #[serde(rename = "canceled")] + Canceled, + #[serde(rename = "expired")] + Expired, +} + #[derive(Deserialize)] struct ResponseBlock { text: Option, @@ -130,9 +172,20 @@ impl ClaudeProvider { model: model.into(), base_url: LOCAL_BASE_URL.to_string(), max_retries, + poll_interval: DEFAULT_BATCH_POLL_INTERVAL, } } + #[cfg(test)] + pub fn with_poll_interval( + self, + interval: std::time::Duration, + ) -> Self { + let mut this = self; + this.poll_interval = interval; + this + } + pub fn with_base_url(self, url: impl Into) -> Self { #[cfg(test)] { @@ -224,32 +277,235 @@ impl ClaudeProvider { .await .context("Failed to parse Claude API response")?; - let truncated = - api_response.stop_reason.as_deref() == Some("max_tokens"); + return response_text(api_response); + } + + Err(last_err.unwrap_or_else(|| { + anyhow::anyhow!("All retry attempts exhausted") + })) + } + + /// Build the Messages API request for a describe payload — shared + /// by the individual and batch paths. + fn describe_api_request( + &self, + payload: &DescribePayload, + context: &DescribeContext, + ) -> ApiRequest { + let system_text = format!( + "{task}\n\n{instructions}", + task = super::describe_system_prompt(), + instructions = super::describe_response_instructions(), + ); + + let content = match payload { + DescribePayload::Image { data, mime_type } => { + use base64::Engine; + let encoded = + base64::engine::general_purpose::STANDARD.encode(data); + vec![ + ContentBlock::Image { + source: ImageSource { + source_type: "base64", + media_type: mime_type.clone(), + data: encoded, + }, + }, + ContentBlock::Text { + text: super::describe_user_prompt(context), + cache_control: None, + }, + ] + } + DescribePayload::Text { excerpt } => vec![ContentBlock::Text { + text: super::describe_text_user_prompt(context, excerpt), + cache_control: None, + }], + }; + + cached_api_request( + self.model.clone(), + 1024, + Some(vec![cached_system_block(system_text)]), + vec![Message { + role: "user", + content, + }], + ) + } + + /// Submit describe requests as a message batch, poll until it + /// ends, and collect per-item results in submission order. + async fn run_message_batch( + &self, + requests: &[DescribeRequest], + ) -> Result>> { + let items: Vec = requests + .iter() + .enumerate() + .map(|(i, r)| BatchRequestItem { + custom_id: format!("req-{i}"), + params: self.describe_api_request(&r.payload, &r.context), + }) + .collect(); + + let endpoint = format!("{}/v1/messages/batches", self.base_url); + let response = self + .client + .post(&endpoint) + .header("x-api-key", &self.api_key) + .header("anthropic-version", "2023-06-01") + .header("content-type", "application/json") + .json(&BatchSubmitBody { requests: items }) + .send() + .await + .context("Failed to submit message batch")?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!( + "Batch submission failed ({}): {}", + status, + preview(&body, 200) + ); + } + + let mut batch: BatchStatus = response + .json() + .await + .context("Failed to parse batch submission response")?; + + tracing::info!( + batch_id = %batch.id, + requests = requests.len(), + "Submitted message batch" + ); + + let status_endpoint = + format!("{}/v1/messages/batches/{}", self.base_url, batch.id); - let raw = api_response - .content - .into_iter() - .find_map(|block| block.text) - .context("No text content in Claude API response")?; + while batch.processing_status != "ended" { + tokio::time::sleep(self.poll_interval).await; + let response = self + .client + .get(&status_endpoint) + .header("x-api-key", &self.api_key) + .header("anthropic-version", "2023-06-01") + .send() + .await + .context("Failed to poll batch status")?; - if truncated { + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); anyhow::bail!( - "Claude response truncated (hit max_tokens limit). \ - Increase max_tokens or reduce the input size. \ - Partial response ({} bytes): {}", - raw.len(), - preview(&raw, 500), + "Batch status poll failed ({}): {}", + status, + preview(&body, 200) ); } - return Ok(extract_json(&raw)); + batch = response + .json() + .await + .context("Failed to parse batch status response")?; + tracing::debug!( + batch_id = %batch.id, + status = %batch.processing_status, + "Polled message batch" + ); } - Err(last_err.unwrap_or_else(|| { - anyhow::anyhow!("All retry attempts exhausted") - })) + let results_endpoint = format!("{status_endpoint}/results"); + let response = self + .client + .get(&results_endpoint) + .header("x-api-key", &self.api_key) + .header("anthropic-version", "2023-06-01") + .send() + .await + .context("Failed to fetch batch results")?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!( + "Batch results fetch failed ({}): {}", + status, + preview(&body, 200) + ); + } + + let body = response + .text() + .await + .context("Failed to read batch results body")?; + + let mut by_id = std::collections::HashMap::new(); + for line in body.lines().filter(|l| !l.trim().is_empty()) { + let parsed: BatchResultLine = serde_json::from_str(line) + .with_context(|| { + format!( + "Failed to parse batch result line: {}", + preview(line, 200) + ) + })?; + + let outcome = match parsed.result { + BatchResult::Succeeded { message } => response_text(message) + .and_then(|text| { + serde_json::from_str::(&text) + .context("Failed to parse description JSON from Claude") + }), + BatchResult::Errored { error } => { + Err(anyhow::anyhow!("Batch item failed: {error}")) + } + BatchResult::Canceled => { + Err(anyhow::anyhow!("Batch item canceled")) + } + BatchResult::Expired => { + Err(anyhow::anyhow!("Batch item expired")) + } + }; + by_id.insert(parsed.custom_id, outcome); + } + + Ok( + (0..requests.len()) + .map(|i| { + by_id.remove(&format!("req-{i}")).unwrap_or_else(|| { + Err(anyhow::anyhow!("Missing batch result for req-{i}")) + }) + }) + .collect(), + ) + } +} + +/// Extract the JSON text payload from a successful API response, +/// rejecting truncated responses. +fn response_text(api_response: ApiResponse) -> Result { + let truncated = + api_response.stop_reason.as_deref() == Some("max_tokens"); + + let raw = api_response + .content + .into_iter() + .find_map(|block| block.text) + .context("No text content in Claude API response")?; + + if truncated { + anyhow::bail!( + "Claude response truncated (hit max_tokens limit). \ + Increase max_tokens or reduce the input size. \ + Partial response ({} bytes): {}", + raw.len(), + preview(&raw, 500), + ); } + + Ok(extract_json(&raw)) } fn is_retryable_status(status: reqwest::StatusCode) -> bool { @@ -341,37 +597,11 @@ impl AiProvider for ClaudeProvider { mime_type: &str, context: &DescribeContext, ) -> Result { - use base64::Engine; - let encoded = - base64::engine::general_purpose::STANDARD.encode(image_data); - let user_text = super::describe_user_prompt(context); - let system_text = format!( - "{task}\n\n{instructions}", - task = super::describe_system_prompt(), - instructions = super::describe_response_instructions(), - ); - - let request = cached_api_request( - self.model.clone(), - 1024, - Some(vec![cached_system_block(system_text)]), - vec![Message { - role: "user", - content: vec![ - ContentBlock::Image { - source: ImageSource { - source_type: "base64", - media_type: mime_type.to_string(), - data: encoded, - }, - }, - ContentBlock::Text { - text: user_text, - cache_control: None, - }, - ], - }], - ); + let payload = DescribePayload::Image { + data: image_data.to_vec(), + mime_type: mime_type.to_string(), + }; + let request = self.describe_api_request(&payload, context); let text = self.send_request(request).await?; let description: ContentDescription = serde_json::from_str(&text) @@ -384,26 +614,10 @@ impl AiProvider for ClaudeProvider { excerpt: &str, context: &DescribeContext, ) -> Result { - let user_text = - super::describe_text_user_prompt(context, excerpt); - let system_text = format!( - "{task}\n\n{instructions}", - task = super::describe_system_prompt(), - instructions = super::describe_response_instructions(), - ); - - let request = cached_api_request( - self.model.clone(), - 1024, - Some(vec![cached_system_block(system_text)]), - vec![Message { - role: "user", - content: vec![ContentBlock::Text { - text: user_text, - cache_control: None, - }], - }], - ); + let payload = DescribePayload::Text { + excerpt: excerpt.to_string(), + }; + let request = self.describe_api_request(&payload, context); let text = self.send_request(request).await?; let description: ContentDescription = serde_json::from_str(&text) @@ -411,6 +625,24 @@ impl AiProvider for ClaudeProvider { Ok(description) } + async fn describe_batch( + &self, + requests: Vec, + ) -> Vec> { + match self.run_message_batch(&requests).await { + Ok(results) => results, + Err(err) => { + // Batch-level failure (submission/poll/fetch): every item + // fails with the shared cause. + let msg = format!("{err:#}"); + requests + .iter() + .map(|_| Err(anyhow::anyhow!("{msg}"))) + .collect() + } + } + } + async fn propose_groups( &self, files: &[FileSummary], @@ -722,4 +954,244 @@ mod tests { let input = " \n {\"key\": \"value\"} \n "; assert_eq!(extract_json(input), r#"{"key": "value"}"#); } + + fn batch_test_requests() -> Vec { + vec![ + DescribeRequest { + payload: DescribePayload::Image { + data: vec![0xFF, 0xD8], + mime_type: "image/jpeg".to_string(), + }, + context: DescribeContext { + filename: "a.jpg".to_string(), + file_type_label: "JPEG image".to_string(), + file_size: 2, + metadata_hint: None, + }, + }, + DescribeRequest { + payload: DescribePayload::Text { + excerpt: "LEASE AGREEMENT".to_string(), + }, + context: DescribeContext { + filename: "lease.txt".to_string(), + file_type_label: "TXT document".to_string(), + file_size: 15, + metadata_hint: None, + }, + }, + ] + } + + fn batch_description_json(summary: &str) -> String { + serde_json::json!({ + "summary": summary, + "tags": ["t"], + "suggested_category": "other", + "confidence": 0.9, + }) + .to_string() + } + + fn batch_result_line(custom_id: &str, summary: &str) -> String { + serde_json::json!({ + "custom_id": custom_id, + "result": { + "type": "succeeded", + "message": { + "content": [ + {"type": "text", "text": batch_description_json(summary)} + ], + "stop_reason": "end_turn", + }, + }, + }) + .to_string() + } + + #[tokio::test] + async fn describe_batch_submits_and_orders_results() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/v1/messages/batches")) + .and(header("x-api-key", "test-key")) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({ + "id": "msgbatch_01", + "processing_status": "ended", + }), + )) + .mount(&server) + .await; + + // Results returned out of submission order on purpose. + let results_body = format!( + "{}\n{}\n", + batch_result_line("req-1", "a lease"), + batch_result_line("req-0", "a photo"), + ); + Mock::given(method("GET")) + .and(path("/v1/messages/batches/msgbatch_01/results")) + .respond_with( + ResponseTemplate::new(200).set_body_string(results_body), + ) + .mount(&server) + .await; + + let provider = ClaudeProvider::new( + "test-key".to_string(), + "claude-opus-4-8".to_string(), + 0, + ) + .with_base_url(server.uri()); + + let results = + provider.describe_batch(batch_test_requests()).await; + + assert_eq!(results.len(), 2); + assert_eq!(results[0].as_ref().unwrap().summary, "a photo"); + assert_eq!(results[1].as_ref().unwrap().summary, "a lease"); + } + + #[tokio::test] + async fn describe_batch_polls_until_ended() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/v1/messages/batches")) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({ + "id": "msgbatch_02", + "processing_status": "in_progress", + }), + )) + .mount(&server) + .await; + + // First poll still in progress, then ended. + Mock::given(method("GET")) + .and(path("/v1/messages/batches/msgbatch_02")) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({ + "id": "msgbatch_02", + "processing_status": "in_progress", + }), + )) + .up_to_n_times(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/v1/messages/batches/msgbatch_02")) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({ + "id": "msgbatch_02", + "processing_status": "ended", + }), + )) + .mount(&server) + .await; + + Mock::given(method("GET")) + .and(path("/v1/messages/batches/msgbatch_02/results")) + .respond_with(ResponseTemplate::new(200).set_body_string( + format!( + "{}\n{}\n", + batch_result_line("req-0", "a photo"), + batch_result_line("req-1", "a lease"), + ), + )) + .mount(&server) + .await; + + let provider = ClaudeProvider::new( + "test-key".to_string(), + "claude-opus-4-8".to_string(), + 0, + ) + .with_base_url(server.uri()) + .with_poll_interval(std::time::Duration::ZERO); + + let results = + provider.describe_batch(batch_test_requests()).await; + + assert_eq!(results.len(), 2); + assert!(results.iter().all(|r| r.is_ok())); + } + + #[tokio::test] + async fn describe_batch_reports_errored_items_individually() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/v1/messages/batches")) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({ + "id": "msgbatch_03", + "processing_status": "ended", + }), + )) + .mount(&server) + .await; + + let errored = serde_json::json!({ + "custom_id": "req-1", + "result": { + "type": "errored", + "error": {"type": "invalid_request", "message": "too large"}, + }, + }) + .to_string(); + Mock::given(method("GET")) + .and(path("/v1/messages/batches/msgbatch_03/results")) + .respond_with(ResponseTemplate::new(200).set_body_string( + format!( + "{}\n{}\n", + batch_result_line("req-0", "a photo"), + errored, + ), + )) + .mount(&server) + .await; + + let provider = ClaudeProvider::new( + "test-key".to_string(), + "claude-opus-4-8".to_string(), + 0, + ) + .with_base_url(server.uri()); + + let results = + provider.describe_batch(batch_test_requests()).await; + + assert!(results[0].is_ok()); + let err = results[1].as_ref().unwrap_err().to_string(); + assert!(err.contains("too large")); + } + + #[tokio::test] + async fn describe_batch_fails_all_items_on_submission_error() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/v1/messages/batches")) + .respond_with( + ResponseTemplate::new(401).set_body_string("bad key"), + ) + .mount(&server) + .await; + + let provider = ClaudeProvider::new( + "test-key".to_string(), + "claude-opus-4-8".to_string(), + 0, + ) + .with_base_url(server.uri()); + + let results = + provider.describe_batch(batch_test_requests()).await; + + assert_eq!(results.len(), 2); + assert!(results.iter().all(|r| r.is_err())); + } } diff --git a/src/ai/mod.rs b/src/ai/mod.rs index ecff90a..1e220bd 100644 --- a/src/ai/mod.rs +++ b/src/ai/mod.rs @@ -15,6 +15,18 @@ pub struct DescribeContext { pub metadata_hint: Option, } +/// Content payload for a single describe request. +pub enum DescribePayload { + Image { data: Vec, mime_type: String }, + Text { excerpt: String }, +} + +/// A self-contained describe request, suitable for batch submission. +pub struct DescribeRequest { + pub payload: DescribePayload, + pub context: DescribeContext, +} + pub trait AiProvider: Send + Sync { fn describe_image( &self, @@ -40,6 +52,33 @@ pub trait AiProvider: Send + Sync { &self, files: &[FileSummary], ) -> impl std::future::Future>> + Send; + + /// Describe many files in one operation. The default falls back to + /// sequential individual calls; providers with a native batch API + /// (50% discount) should override this. + fn describe_batch( + &self, + requests: Vec, + ) -> impl std::future::Future>> + + Send { + async move { + let mut results = Vec::with_capacity(requests.len()); + for request in &requests { + let result = match &request.payload { + DescribePayload::Image { data, mime_type } => { + self + .describe_image(data, mime_type, &request.context) + .await + } + DescribePayload::Text { excerpt } => { + self.describe_text(excerpt, &request.context).await + } + }; + results.push(result); + } + results + } + } } #[cfg(test)] @@ -59,4 +98,95 @@ mod tests { assert_eq!(ctx.file_size, 2048); assert!(ctx.metadata_hint.is_some()); } + + #[tokio::test] + async fn default_describe_batch_falls_back_to_individual_calls() { + struct CountingProvider { + images: std::sync::atomic::AtomicUsize, + texts: std::sync::atomic::AtomicUsize, + } + + impl AiProvider for CountingProvider { + async fn describe_image( + &self, + _: &[u8], + _: &str, + _: &DescribeContext, + ) -> Result { + self + .images + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(stub_description("an image")) + } + + async fn describe_text( + &self, + _: &str, + _: &DescribeContext, + ) -> Result { + self.texts.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(stub_description("a text")) + } + + async fn propose_groups( + &self, + _: &[FileSummary], + ) -> Result> { + panic!("unused"); + } + } + + fn stub_description(summary: &str) -> ContentDescription { + ContentDescription { + summary: summary.to_string(), + tags: vec![], + suggested_category: "other".to_string(), + confidence: 0.5, + } + } + + fn ctx(name: &str) -> DescribeContext { + DescribeContext { + filename: name.to_string(), + file_type_label: "test".to_string(), + file_size: 1, + metadata_hint: None, + } + } + + let provider = CountingProvider { + images: std::sync::atomic::AtomicUsize::new(0), + texts: std::sync::atomic::AtomicUsize::new(0), + }; + + let requests = vec![ + DescribeRequest { + payload: DescribePayload::Image { + data: vec![0xFF], + mime_type: "image/jpeg".to_string(), + }, + context: ctx("a.jpg"), + }, + DescribeRequest { + payload: DescribePayload::Text { + excerpt: "hello".to_string(), + }, + context: ctx("b.txt"), + }, + ]; + + let results = provider.describe_batch(requests).await; + + assert_eq!(results.len(), 2); + assert_eq!(results[0].as_ref().unwrap().summary, "an image"); + assert_eq!(results[1].as_ref().unwrap().summary, "a text"); + assert_eq!( + provider.images.load(std::sync::atomic::Ordering::SeqCst), + 1 + ); + assert_eq!( + provider.texts.load(std::sync::atomic::Ordering::SeqCst), + 1 + ); + } } diff --git a/src/analyze/mod.rs b/src/analyze/mod.rs index 644f213..6c1a1e7 100644 --- a/src/analyze/mod.rs +++ b/src/analyze/mod.rs @@ -9,6 +9,9 @@ use crate::model::{ContentDescription, FingerprintedFile}; pub struct AnalyzeOptions { pub cache_dir: PathBuf, pub max_concurrent: usize, + /// Use the Batch API (50% cheaper, async) instead of concurrent + /// individual requests. + pub use_batch_api: bool, } impl Default for AnalyzeOptions { @@ -16,6 +19,7 @@ impl Default for AnalyzeOptions { Self { cache_dir: default_cache_dir(), max_concurrent: 5, + use_batch_api: false, } } } @@ -368,6 +372,10 @@ pub async fn analyze_batch( files: &[FingerprintedFile], options: &AnalyzeOptions, ) -> Vec> { + if options.use_batch_api { + return analyze_batch_via_api(provider, files, options).await; + } + use futures::stream::{self, StreamExt}; let semaphore = @@ -389,6 +397,145 @@ pub async fn analyze_batch( .await } +/// What a file contributes to a batch run. +enum BatchSlot { + /// Already resolved locally (cache hit, fallback description, or + /// a local error such as an unreadable image). + Resolved(Result), + /// Needs an API call; index into the submitted request list. + Submitted(usize), +} + +/// Analyze files through the provider's batch interface. Cache hits +/// and filename-only fallbacks resolve locally; everything else is +/// submitted as one batch. Videos go through the regular per-file +/// path since keyframe extraction is multi-request. +async fn analyze_batch_via_api( + provider: &impl AiProvider, + files: &[FingerprintedFile], + options: &AnalyzeOptions, +) -> Vec> { + use crate::ai::{DescribePayload, DescribeRequest}; + + let mut slots = Vec::with_capacity(files.len()); + let mut requests = Vec::new(); + + for file in files { + if let Some(cached) = + read_cache(&options.cache_dir, &file.blake3_hash).await + { + slots.push(BatchSlot::Resolved(Ok(cached))); + continue; + } + + let filename = file + .scanned + .path + .file_name() + .unwrap_or_default() + .to_string_lossy() + .to_string(); + + if file.scanned.file_type.is_video() { + slots.push(BatchSlot::Resolved( + analyze_file(provider, file, options).await, + )); + continue; + } + + if file.scanned.file_type.is_image() { + match tokio::fs::read(&file.scanned.path).await { + Ok(data) => { + requests.push(DescribeRequest { + payload: DescribePayload::Image { + data, + mime_type: file + .scanned + .file_type + .mime_type() + .to_string(), + }, + context: DescribeContext { + filename, + file_type_label: file + .scanned + .file_type + .mime_type() + .to_string(), + file_size: file.scanned.size, + metadata_hint: None, + }, + }); + slots.push(BatchSlot::Submitted(requests.len() - 1)); + } + Err(e) => { + slots.push(BatchSlot::Resolved(Err(anyhow::anyhow!( + "Failed to read file {}: {e}", + file.scanned.path.display() + )))); + } + } + continue; + } + + if matches!( + file.scanned.file_type, + crate::model::FileType::Document(_) + ) { + match extract_document_text(file).await { + Some(text) if !text.trim().is_empty() => { + requests.push(DescribeRequest { + payload: DescribePayload::Text { excerpt: text }, + context: DescribeContext { + filename, + file_type_label: document_type_label(file), + file_size: file.scanned.size, + metadata_hint: None, + }, + }); + slots.push(BatchSlot::Submitted(requests.len() - 1)); + } + _ => { + slots.push(BatchSlot::Resolved(Ok(describe_by_filename( + file, &filename, + )))); + } + } + continue; + } + + slots.push(BatchSlot::Resolved(Ok(describe_by_filename( + file, &filename, + )))); + } + + let mut batch_results = if requests.is_empty() { + Vec::new() + } else { + provider.describe_batch(requests).await + }; + + let mut out = Vec::with_capacity(files.len()); + for (file, slot) in files.iter().zip(slots) { + let result = match slot { + BatchSlot::Resolved(r) => r, + BatchSlot::Submitted(i) => std::mem::replace( + &mut batch_results[i], + Err(anyhow::anyhow!("Batch result already taken")), + ), + }; + + if let Ok(desc) = &result { + let _ = + write_cache(&options.cache_dir, &file.blake3_hash, desc) + .await; + } + out.push(result); + } + + out +} + #[cfg(test)] mod tests { use super::*; @@ -514,6 +661,7 @@ mod tests { let opts = AnalyzeOptions { cache_dir: dir.path().to_path_buf(), max_concurrent: 1, + use_batch_api: false, }; struct PanicProvider; @@ -550,6 +698,7 @@ mod tests { let opts = AnalyzeOptions { cache_dir: cache_dir.path().to_path_buf(), max_concurrent: 1, + use_batch_api: false, }; struct FakeProvider; @@ -591,6 +740,7 @@ mod tests { let opts = AnalyzeOptions { cache_dir: cache_dir.path().to_path_buf(), max_concurrent: 1, + use_batch_api: false, }; struct FakeProvider; @@ -651,6 +801,7 @@ mod tests { let opts = AnalyzeOptions { cache_dir: cache_dir.path().to_path_buf(), max_concurrent: 1, + use_batch_api: false, }; struct UnusedProvider; @@ -691,6 +842,7 @@ mod tests { let opts = AnalyzeOptions { cache_dir: cache_dir.path().to_path_buf(), max_concurrent: 1, + use_batch_api: false, }; struct StubProvider; @@ -785,6 +937,7 @@ mod tests { let opts = AnalyzeOptions { cache_dir: cache_dir.path().to_path_buf(), max_concurrent: 1, + use_batch_api: false, }; let provider = TextCapturingProvider { @@ -815,6 +968,7 @@ mod tests { let opts = AnalyzeOptions { cache_dir: cache_dir.path().to_path_buf(), max_concurrent: 1, + use_batch_api: false, }; let provider = TextCapturingProvider { @@ -842,6 +996,7 @@ mod tests { let opts = AnalyzeOptions { cache_dir: cache_dir.path().to_path_buf(), max_concurrent: 1, + use_batch_api: false, }; let provider = TextCapturingProvider { @@ -877,4 +1032,170 @@ mod tests { assert!(truncated.len() <= MAX_TEXT_EXCERPT_BYTES); assert!(truncated.is_char_boundary(truncated.len())); } + + /// Provider that only answers through describe_batch and records + /// how many requests it received. + struct BatchOnlyProvider { + batch_sizes: std::sync::Mutex>, + } + + impl BatchOnlyProvider { + fn new() -> Self { + Self { + batch_sizes: std::sync::Mutex::new(Vec::new()), + } + } + } + + impl AiProvider for BatchOnlyProvider { + async fn describe_image( + &self, + _: &[u8], + _: &str, + _: &DescribeContext, + ) -> Result { + panic!("individual describe_image used in batch mode"); + } + + async fn describe_text( + &self, + _: &str, + _: &DescribeContext, + ) -> Result { + panic!("individual describe_text used in batch mode"); + } + + async fn describe_batch( + &self, + requests: Vec, + ) -> Vec> { + self.batch_sizes.lock().unwrap().push(requests.len()); + requests + .iter() + .map(|r| { + Ok(ContentDescription { + summary: format!("batched: {}", r.context.filename), + tags: vec!["batch".to_string()], + suggested_category: "other".to_string(), + confidence: 0.9, + }) + }) + .collect() + } + + async fn propose_groups( + &self, + _: &[crate::model::FileSummary], + ) -> Result> { + panic!("unused"); + } + } + + #[tokio::test] + async fn analyze_batch_uses_batch_api_when_enabled() { + let cache_dir = TempDir::new().unwrap(); + let file_dir = TempDir::new().unwrap(); + let image = + make_test_file(file_dir.path(), "photo.jpg", b"jpeg bytes"); + let doc = make_document_file( + file_dir.path(), + "lease.txt", + b"LEASE AGREEMENT terms", + crate::model::DocumentFormat::Txt, + ); + + let opts = AnalyzeOptions { + cache_dir: cache_dir.path().to_path_buf(), + max_concurrent: 1, + use_batch_api: true, + }; + + let provider = BatchOnlyProvider::new(); + let results = + analyze_batch(&provider, &[image.clone(), doc], &opts).await; + + assert_eq!(results.len(), 2); + assert_eq!( + results[0].as_ref().unwrap().summary, + "batched: photo.jpg" + ); + assert_eq!( + results[1].as_ref().unwrap().summary, + "batched: lease.txt" + ); + // One batch containing both requests + assert_eq!(*provider.batch_sizes.lock().unwrap(), vec![2]); + + // Results were cached for the next run + assert!(read_cache(cache_dir.path(), &image.blake3_hash) + .await + .is_some()); + } + + #[tokio::test] + async fn analyze_batch_api_skips_cached_files() { + let cache_dir = TempDir::new().unwrap(); + let file_dir = TempDir::new().unwrap(); + let cached_file = + make_test_file(file_dir.path(), "seen.jpg", b"old bytes"); + let new_file = + make_test_file(file_dir.path(), "new.jpg", b"new bytes"); + + write_cache( + cache_dir.path(), + &cached_file.blake3_hash, + &sample_description(), + ) + .await + .unwrap(); + + let opts = AnalyzeOptions { + cache_dir: cache_dir.path().to_path_buf(), + max_concurrent: 1, + use_batch_api: true, + }; + + let provider = BatchOnlyProvider::new(); + let results = + analyze_batch(&provider, &[cached_file, new_file], &opts).await; + + assert_eq!( + results[0].as_ref().unwrap().summary, + "A sunset over the ocean" + ); + assert_eq!( + results[1].as_ref().unwrap().summary, + "batched: new.jpg" + ); + // Only the uncached file was submitted + assert_eq!(*provider.batch_sizes.lock().unwrap(), vec![1]); + } + + #[tokio::test] + async fn analyze_batch_api_resolves_unbatchable_files_locally() { + let cache_dir = TempDir::new().unwrap(); + let file_dir = TempDir::new().unwrap(); + let docx = make_document_file( + file_dir.path(), + "report.docx", + b"PK\x03\x04 binary", + crate::model::DocumentFormat::Docx, + ); + + let opts = AnalyzeOptions { + cache_dir: cache_dir.path().to_path_buf(), + max_concurrent: 1, + use_batch_api: true, + }; + + let provider = BatchOnlyProvider::new(); + let results = analyze_batch(&provider, &[docx], &opts).await; + + assert_eq!( + results[0].as_ref().unwrap().suggested_category, + "document" + ); + // Nothing was submitted to the API + assert!(provider.batch_sizes.lock().unwrap().is_empty()); + } } diff --git a/src/config.rs b/src/config.rs index 861b595..a462e6b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -40,6 +40,11 @@ pub struct CliArgs { #[arg(long)] pub max_cost: Option, + /// Use the Batch API for analysis (50% cheaper, results may take + /// minutes to hours) + #[arg(long)] + pub batch: bool, + /// Anthropic API base URL (for proxies like whetstone) #[arg(long, env = "ANTHROPIC_BASE_URL")] pub api_base_url: Option, @@ -280,6 +285,7 @@ mod tests { api_key: None, max_files: None, max_cost: None, + batch: false, api_base_url: None, verbose: 0, include_trash: false, @@ -585,4 +591,16 @@ base_url = "https://from-file.example.com" assert!(ai.base_url.is_none()); } + + #[test] + fn batch_flag_defaults_off_and_parses() { + use clap::Parser; + + let default_args = CliArgs::parse_from(["spindle", "some-dir"]); + assert!(!default_args.batch); + + let batch_args = + CliArgs::parse_from(["spindle", "--batch", "some-dir"]); + assert!(batch_args.batch); + } } diff --git a/src/main.rs b/src/main.rs index 70e7ba1..a7ae131 100644 --- a/src/main.rs +++ b/src/main.rs @@ -62,6 +62,7 @@ async fn main() -> Result<()> { max_concurrent: config.ai.max_concurrent_requests, include_trash: cli.include_trash, type_filter: cli.file_types.clone(), + use_batch_api: cli.batch, }; let (tx, mut rx) = tokio::sync::mpsc::channel::(64); diff --git a/src/pipeline.rs b/src/pipeline.rs index 162cbe4..dec2b4a 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -74,6 +74,7 @@ pub struct PipelineConfig { pub max_concurrent: usize, pub include_trash: bool, pub type_filter: Vec, + pub use_batch_api: bool, } #[derive(Debug)] @@ -248,6 +249,7 @@ async fn run_ai_pipeline( let opts = AnalyzeOptions { cache_dir: config.cache_dir.clone(), max_concurrent: config.max_concurrent, + use_batch_api: config.use_batch_api, }; let subset: Vec<_> = @@ -404,6 +406,7 @@ mod tests { max_concurrent: 2, include_trash: false, type_filter: vec![], + use_batch_api: false, }; let (tx, mut rx) = mpsc::channel(64); @@ -453,6 +456,7 @@ mod tests { max_concurrent: 2, include_trash: false, type_filter: vec![], + use_batch_api: false, }; let (tx, mut rx) = mpsc::channel(64); @@ -493,6 +497,7 @@ mod tests { max_concurrent: 2, include_trash: false, type_filter: vec![], + use_batch_api: false, }; let (tx, _rx) = mpsc::channel(64); @@ -531,6 +536,7 @@ mod tests { max_concurrent: 2, include_trash: false, type_filter: vec![], + use_batch_api: false, }; let (tx, _rx) = mpsc::channel(64); @@ -565,6 +571,7 @@ mod tests { max_concurrent: 2, include_trash: false, type_filter: vec![], + use_batch_api: false, }; let (tx, mut rx) = mpsc::channel(64); diff --git a/tests/integration_test.rs b/tests/integration_test.rs index f40e350..5c5799e 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -276,6 +276,7 @@ async fn analyze_caches_and_reuses_results() { let opts = AnalyzeOptions { cache_dir: cache_dir.path().to_path_buf(), max_concurrent: 2, + use_batch_api: false, }; let result1 = @@ -321,6 +322,7 @@ async fn analyze_batch_processes_multiple_files() { let opts = AnalyzeOptions { cache_dir: cache_dir.path().to_path_buf(), max_concurrent: 2, + use_batch_api: false, }; let results = analyze_batch(&FakeAiProvider, &files, &opts).await; @@ -360,6 +362,7 @@ async fn full_pipeline_end_to_end() { let opts = AnalyzeOptions { cache_dir: cache_dir.path().to_path_buf(), max_concurrent: 2, + use_batch_api: false, }; let descriptions = analyze_batch(&FakeAiProvider, &fingerprinted, &opts).await;