use anthropic::{ ANTHROPIC_API_URL, Event, Message, Request as AnthropicRequest, RequestContent, Response as AnthropicResponse, ResponseContent, Role, non_streaming_completion, stream_completion, }; use anyhow::Result; use futures::StreamExt as _; use http_client::HttpClient; use indoc::indoc; use reqwest_client::ReqwestClient; use sqlez::bindable::Bind; use sqlez::bindable::StaticColumnCount; use sqlez_macros::sql; use std::collections::HashSet; use std::hash::Hash; use std::hash::Hasher; use std::path::Path; use std::sync::{Arc, Mutex}; pub struct PlainLlmClient { pub http_client: Arc, pub api_key: String, } impl PlainLlmClient { pub fn new() -> Result { let http_client: Arc = Arc::new(ReqwestClient::new()); let api_key = std::env::var("ANTHROPIC_API_KEY") .map_err(|_| anyhow::anyhow!("ANTHROPIC_API_KEY environment variable not set"))?; Ok(Self { http_client, api_key, }) } pub async fn generate( &self, model: &str, max_tokens: u64, messages: Vec, ) -> Result { let request = AnthropicRequest { model: model.to_string(), max_tokens, messages, tools: Vec::new(), thinking: None, tool_choice: None, system: None, cache_control: None, context_management: None, metadata: None, output_config: None, stop_sequences: Vec::new(), speed: None, temperature: None, top_k: None, top_p: None, }; let response = non_streaming_completion( self.http_client.as_ref(), ANTHROPIC_API_URL, &self.api_key, request, None, &http_client::CustomHeaders::default(), ) .await .map_err(|e| anyhow::anyhow!("{:?}", e))?; Ok(response) } pub async fn generate_streaming( &self, model: &str, max_tokens: u64, messages: Vec, mut on_progress: F, ) -> Result where F: FnMut(usize, &str), { let request = AnthropicRequest { model: model.to_string(), max_tokens, messages, tools: Vec::new(), thinking: None, tool_choice: None, system: None, cache_control: None, context_management: None, metadata: None, output_config: None, stop_sequences: Vec::new(), speed: None, temperature: None, top_k: None, top_p: None, }; let mut stream = stream_completion( self.http_client.as_ref(), ANTHROPIC_API_URL, &self.api_key, request, None, &http_client::CustomHeaders::default(), ) .await .map_err(|e| anyhow::anyhow!("{:?}", e))?; let mut response: Option = None; let mut text_content = String::new(); while let Some(event_result) = stream.next().await { let event = event_result.map_err(|e| anyhow::anyhow!("{:?}", e))?; match event { Event::MessageStart { message } => { response = Some(message); } Event::ContentBlockDelta { delta, .. } => { if let anthropic::ContentDelta::TextDelta { text } = delta { text_content.push_str(&text); on_progress(text_content.len(), &text_content); } } _ => {} } } let mut response = response.ok_or_else(|| anyhow::anyhow!("No response received"))?; if response.content.is_empty() && !text_content.is_empty() { response .content .push(ResponseContent::Text { text: text_content }); } Ok(response) } } pub struct BatchingLlmClient { connection: Mutex, http_client: Arc, api_key: String, } struct CacheRow { request_hash: String, request: Option, response: Option, batch_id: Option, } impl StaticColumnCount for CacheRow { fn column_count() -> usize { 4 } } impl Bind for CacheRow { fn bind(&self, statement: &sqlez::statement::Statement, start_index: i32) -> Result { let next_index = statement.bind(&self.request_hash, start_index)?; let next_index = statement.bind(&self.request, next_index)?; let next_index = statement.bind(&self.response, next_index)?; let next_index = statement.bind(&self.batch_id, next_index)?; Ok(next_index) } } #[derive(serde::Serialize, serde::Deserialize)] struct SerializableRequest { model: String, max_tokens: u64, messages: Vec, } #[derive(serde::Serialize, serde::Deserialize)] struct SerializableMessage { role: String, content: String, } impl BatchingLlmClient { fn new(cache_path: &Path) -> Result { let http_client: Arc = Arc::new(ReqwestClient::new()); let api_key = std::env::var("ANTHROPIC_API_KEY") .map_err(|_| anyhow::anyhow!("ANTHROPIC_API_KEY environment variable not set"))?; let connection = sqlez::connection::Connection::open_file(&cache_path.to_str().unwrap()); let mut statement = sqlez::statement::Statement::prepare( &connection, indoc! {" CREATE TABLE IF NOT EXISTS cache ( request_hash TEXT PRIMARY KEY, request TEXT, response TEXT, batch_id TEXT ); "}, )?; statement.exec()?; drop(statement); Ok(Self { connection: Mutex::new(connection), http_client, api_key, }) } pub fn lookup( &self, model: &str, max_tokens: u64, messages: &[Message], seed: Option, ) -> Result> { let request_hash_str = Self::request_hash(model, max_tokens, messages, seed); let connection = self.connection.lock().unwrap(); let response: Vec = connection.select_bound( &sql!(SELECT response FROM cache WHERE request_hash = ?1 AND response IS NOT NULL;), )?(request_hash_str.as_str())?; Ok(response .into_iter() .next() .and_then(|text| serde_json::from_str(&text).ok())) } pub fn mark_for_batch( &self, model: &str, max_tokens: u64, messages: &[Message], seed: Option, ) -> Result<()> { let request_hash = Self::request_hash(model, max_tokens, messages, seed); let serializable_messages: Vec = messages .iter() .map(|msg| SerializableMessage { role: match msg.role { Role::User => "user".to_string(), Role::Assistant => "assistant".to_string(), }, content: message_content_to_string(&msg.content), }) .collect(); let serializable_request = SerializableRequest { model: model.to_string(), max_tokens, messages: serializable_messages, }; let request = Some(serde_json::to_string(&serializable_request)?); let cache_row = CacheRow { request_hash, request, response: None, batch_id: None, }; let connection = self.connection.lock().unwrap(); connection.exec_bound::(sql!( INSERT OR IGNORE INTO cache(request_hash, request, response, batch_id) VALUES (?, ?, ?, ?)))?( cache_row, ) } async fn generate( &self, model: &str, max_tokens: u64, messages: Vec, seed: Option, cache_only: bool, ) -> Result> { let response = self.lookup(model, max_tokens, &messages, seed)?; if let Some(response) = response { return Ok(Some(response)); } if !cache_only { self.mark_for_batch(model, max_tokens, &messages, seed)?; } Ok(None) } /// Uploads pending requests as batches (chunked to 16k each); downloads finished batches if any. async fn sync_batches(&self) -> Result<()> { let _batch_ids = self.upload_pending_requests().await?; self.download_finished_batches().await } pub fn pending_batch_count(&self) -> Result { let connection = self.connection.lock().unwrap(); let counts: Vec = connection.select( sql!(SELECT COUNT(*) FROM cache WHERE batch_id IS NOT NULL AND response IS NULL), )?()?; Ok(counts.into_iter().next().unwrap_or(0) as usize) } /// Import batch results from external batch IDs (useful for recovering after database loss) pub async fn import_batches(&self, batch_ids: &[String]) -> Result<()> { for batch_id in batch_ids { log::info!("Importing batch {}", batch_id); let batch_status = anthropic::batches::retrieve_batch( self.http_client.as_ref(), ANTHROPIC_API_URL, &self.api_key, batch_id, ) .await .map_err(|e| anyhow::anyhow!("Failed to retrieve batch {}: {:?}", batch_id, e))?; log::info!( "Batch {} status: {}", batch_id, batch_status.processing_status ); if batch_status.processing_status != "ended" { log::warn!( "Batch {} is not finished (status: {}), skipping", batch_id, batch_status.processing_status ); continue; } let results = anthropic::batches::retrieve_batch_results( self.http_client.as_ref(), ANTHROPIC_API_URL, &self.api_key, batch_id, ) .await .map_err(|e| { anyhow::anyhow!("Failed to retrieve batch results for {}: {:?}", batch_id, e) })?; let mut updates: Vec<(String, String, String)> = Vec::new(); let mut success_count = 0; let mut error_count = 0; for result in results { let request_hash = result .custom_id .strip_prefix("req_hash_") .unwrap_or(&result.custom_id) .to_string(); match result.result { anthropic::batches::BatchResult::Succeeded { message } => { let response_json = serde_json::to_string(&message)?; updates.push((request_hash, response_json, batch_id.clone())); success_count += 1; } anthropic::batches::BatchResult::Errored { error } => { log::error!( "Batch request {} failed: {}: {}", request_hash, error.error.error_type, error.error.message ); let error_json = serde_json::json!({ "error": { "type": error.error.error_type, "message": error.error.message } }) .to_string(); updates.push((request_hash, error_json, batch_id.clone())); error_count += 1; } anthropic::batches::BatchResult::Canceled => { log::warn!("Batch request {} was canceled", request_hash); error_count += 1; } anthropic::batches::BatchResult::Expired => { log::warn!("Batch request {} expired", request_hash); error_count += 1; } } } let connection = self.connection.lock().unwrap(); connection.with_savepoint("batch_import", || { // Use INSERT OR REPLACE to handle both new entries and updating existing ones let q = sql!( INSERT OR REPLACE INTO cache(request_hash, request, response, batch_id) VALUES (?, (SELECT request FROM cache WHERE request_hash = ?), ?, ?) ); let mut exec = connection.exec_bound::<(&str, &str, &str, &str)>(q)?; for (request_hash, response_json, batch_id) in &updates { exec(( request_hash.as_str(), request_hash.as_str(), response_json.as_str(), batch_id.as_str(), ))?; } Ok(()) })?; log::info!( "Imported batch {}: {} successful, {} errors", batch_id, success_count, error_count ); } Ok(()) } async fn download_finished_batches(&self) -> Result<()> { let batch_ids: Vec = { let connection = self.connection.lock().unwrap(); let q = sql!(SELECT DISTINCT batch_id FROM cache WHERE batch_id IS NOT NULL AND response IS NULL); connection.select(q)?()? }; for batch_id in &batch_ids { let batch_status = anthropic::batches::retrieve_batch( self.http_client.as_ref(), ANTHROPIC_API_URL, &self.api_key, &batch_id, ) .await .map_err(|e| anyhow::anyhow!("{:?}", e))?; log::info!( "Batch {} status: {}", batch_id, batch_status.processing_status ); if batch_status.processing_status == "ended" { let results = anthropic::batches::retrieve_batch_results( self.http_client.as_ref(), ANTHROPIC_API_URL, &self.api_key, &batch_id, ) .await .map_err(|e| anyhow::anyhow!("{:?}", e))?; let mut updates: Vec<(String, String)> = Vec::new(); let mut success_count = 0; for result in results { let request_hash = result .custom_id .strip_prefix("req_hash_") .unwrap_or(&result.custom_id) .to_string(); match result.result { anthropic::batches::BatchResult::Succeeded { message } => { let response_json = serde_json::to_string(&message)?; updates.push((response_json, request_hash)); success_count += 1; } anthropic::batches::BatchResult::Errored { error } => { log::error!( "Batch request {} failed: {}: {}", request_hash, error.error.error_type, error.error.message ); let error_json = serde_json::json!({ "error": { "type": error.error.error_type, "message": error.error.message } }) .to_string(); updates.push((error_json, request_hash)); } anthropic::batches::BatchResult::Canceled => { log::warn!("Batch request {} was canceled", request_hash); let error_json = serde_json::json!({ "error": { "type": "canceled", "message": "Batch request was canceled" } }) .to_string(); updates.push((error_json, request_hash)); } anthropic::batches::BatchResult::Expired => { log::warn!("Batch request {} expired", request_hash); let error_json = serde_json::json!({ "error": { "type": "expired", "message": "Batch request expired" } }) .to_string(); updates.push((error_json, request_hash)); } } } let connection = self.connection.lock().unwrap(); connection.with_savepoint("batch_download", || { let q = sql!(UPDATE cache SET response = ? WHERE request_hash = ?); let mut exec = connection.exec_bound::<(&str, &str)>(q)?; for (response_json, request_hash) in &updates { exec((response_json.as_str(), request_hash.as_str()))?; } Ok(()) })?; log::info!("Downloaded {} successful requests", success_count); } } Ok(()) } async fn upload_pending_requests(&self) -> Result> { const BATCH_CHUNK_SIZE: i32 = 16_000; const MAX_BATCH_SIZE_BYTES: usize = 100 * 1024 * 1024; let mut all_batch_ids = Vec::new(); let mut total_uploaded = 0; let mut current_batch_rows = Vec::new(); let mut current_batch_size = 0usize; let mut pending_hashes: HashSet = HashSet::new(); loop { let rows: Vec<(String, String)> = { let connection = self.connection.lock().unwrap(); let q = sql!( SELECT request_hash, request FROM cache WHERE batch_id IS NULL AND response IS NULL LIMIT ? ); connection.select_bound(q)?(BATCH_CHUNK_SIZE)? }; if rows.is_empty() { break; } // Split rows into sub-batches based on size let mut batches_to_upload = Vec::new(); let mut new_rows_added = 0; for row in rows { let (hash, request_str) = row; // Skip rows already added to current_batch_rows but not yet uploaded if pending_hashes.contains(&hash) { continue; } let serializable_request: SerializableRequest = serde_json::from_str(&request_str)?; let messages: Vec = serializable_request .messages .into_iter() .map(|msg| Message { role: match msg.role.as_str() { "user" => Role::User, "assistant" => Role::Assistant, _ => Role::User, }, content: vec![RequestContent::Text { text: msg.content, cache_control: None, }], }) .collect(); let params = AnthropicRequest { model: serializable_request.model, max_tokens: serializable_request.max_tokens, messages, tools: Vec::new(), thinking: None, tool_choice: None, system: None, cache_control: None, context_management: None, metadata: None, output_config: None, stop_sequences: Vec::new(), temperature: None, top_k: None, top_p: None, speed: None, }; let custom_id = format!("req_hash_{}", hash); let batch_request = anthropic::batches::BatchRequest { custom_id, params }; // Estimate the serialized size of this request let estimated_size = serde_json::to_string(&batch_request)?.len(); // If adding this request would exceed the limit, start a new batch if !current_batch_rows.is_empty() && current_batch_size + estimated_size > MAX_BATCH_SIZE_BYTES { batches_to_upload.push((current_batch_rows, current_batch_size)); current_batch_rows = Vec::new(); current_batch_size = 0; } pending_hashes.insert(hash.clone()); current_batch_rows.push((hash, batch_request)); current_batch_size += estimated_size; new_rows_added += 1; } // If no new rows were added this iteration, all pending requests are already // in current_batch_rows, so we should break to avoid an infinite loop if new_rows_added == 0 { break; } // Only upload full batches, keep the partial batch for the next iteration // Upload each sub-batch for (batch_rows, batch_size) in batches_to_upload { let request_hashes: Vec = batch_rows.iter().map(|(hash, _)| hash.clone()).collect(); // Remove uploaded hashes from pending set for hash in &request_hashes { pending_hashes.remove(hash); } let batch_requests: Vec = batch_rows.into_iter().map(|(_, req)| req).collect(); let batch_len = batch_requests.len(); log::info!( "Uploading batch with {} requests (~{:.2} MB)", batch_len, batch_size as f64 / (1024.0 * 1024.0) ); let batch = anthropic::batches::create_batch( self.http_client.as_ref(), ANTHROPIC_API_URL, &self.api_key, anthropic::batches::CreateBatchRequest { requests: batch_requests, }, ) .await .map_err(|e| anyhow::anyhow!("{:?}", e))?; { let connection = self.connection.lock().unwrap(); connection.with_savepoint("batch_upload", || { let q = sql!(UPDATE cache SET batch_id = ? WHERE request_hash = ?); let mut exec = connection.exec_bound::<(&str, &str)>(q)?; for hash in &request_hashes { exec((batch.id.as_str(), hash.as_str()))?; } Ok(()) })?; } total_uploaded += batch_len; log::info!( "Uploaded batch {} with {} requests ({} total)", batch.id, batch_len, total_uploaded ); all_batch_ids.push(batch.id); } } // Upload any remaining partial batch at the end if !current_batch_rows.is_empty() { let request_hashes: Vec = current_batch_rows .iter() .map(|(hash, _)| hash.clone()) .collect(); let batch_requests: Vec = current_batch_rows.into_iter().map(|(_, req)| req).collect(); let batch_len = batch_requests.len(); log::info!( "Uploading final batch with {} requests (~{:.2} MB)", batch_len, current_batch_size as f64 / (1024.0 * 1024.0) ); let batch = anthropic::batches::create_batch( self.http_client.as_ref(), ANTHROPIC_API_URL, &self.api_key, anthropic::batches::CreateBatchRequest { requests: batch_requests, }, ) .await .map_err(|e| anyhow::anyhow!("{:?}", e))?; { let connection = self.connection.lock().unwrap(); connection.with_savepoint("batch_upload", || { let q = sql!(UPDATE cache SET batch_id = ? WHERE request_hash = ?); let mut exec = connection.exec_bound::<(&str, &str)>(q)?; for hash in &request_hashes { exec((batch.id.as_str(), hash.as_str()))?; } Ok(()) })?; } total_uploaded += batch_len; log::info!( "Uploaded batch {} with {} requests ({} total)", batch.id, batch_len, total_uploaded ); all_batch_ids.push(batch.id); } if !all_batch_ids.is_empty() { log::info!( "Finished uploading {} batches with {} total requests", all_batch_ids.len(), total_uploaded ); } Ok(all_batch_ids) } fn request_hash( model: &str, max_tokens: u64, messages: &[Message], seed: Option, ) -> String { let mut hasher = std::hash::DefaultHasher::new(); model.hash(&mut hasher); max_tokens.hash(&mut hasher); for msg in messages { message_content_to_string(&msg.content).hash(&mut hasher); } if let Some(seed) = seed { seed.hash(&mut hasher); } let request_hash = hasher.finish(); format!("{request_hash:016x}") } } fn message_content_to_string(content: &[RequestContent]) -> String { content .iter() .filter_map(|c| match c { RequestContent::Text { text, .. } => Some(text.clone()), _ => None, }) .collect::>() .join("\n") } pub enum AnthropicClient { // No batching Plain(PlainLlmClient), Batch(BatchingLlmClient), Dummy, } impl AnthropicClient { pub fn plain() -> Result { Ok(Self::Plain(PlainLlmClient::new()?)) } pub fn batch(cache_path: &Path) -> Result { Ok(Self::Batch(BatchingLlmClient::new(cache_path)?)) } #[allow(dead_code)] pub fn dummy() -> Self { Self::Dummy } pub async fn generate( &self, model: &str, max_tokens: u64, messages: Vec, seed: Option, cache_only: bool, ) -> Result> { match self { AnthropicClient::Plain(plain_llm_client) => plain_llm_client .generate(model, max_tokens, messages) .await .map(Some), AnthropicClient::Batch(batching_llm_client) => { batching_llm_client .generate(model, max_tokens, messages, seed, cache_only) .await } AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"), } } #[allow(dead_code)] pub async fn generate_streaming( &self, model: &str, max_tokens: u64, messages: Vec, on_progress: F, ) -> Result> where F: FnMut(usize, &str), { match self { AnthropicClient::Plain(plain_llm_client) => plain_llm_client .generate_streaming(model, max_tokens, messages, on_progress) .await .map(Some), AnthropicClient::Batch(_) => { anyhow::bail!("Streaming not supported with batching client") } AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"), } } pub async fn sync_batches(&self) -> Result<()> { match self { AnthropicClient::Plain(_) => Ok(()), AnthropicClient::Batch(batching_llm_client) => batching_llm_client.sync_batches().await, AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"), } } pub fn pending_batch_count(&self) -> Result { match self { AnthropicClient::Plain(_) => Ok(0), AnthropicClient::Batch(batching_llm_client) => { batching_llm_client.pending_batch_count() } AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"), } } pub async fn import_batches(&self, batch_ids: &[String]) -> Result<()> { match self { AnthropicClient::Plain(_) => { anyhow::bail!("Import batches is only supported with batching client") } AnthropicClient::Batch(batching_llm_client) => { batching_llm_client.import_batches(batch_ids).await } AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"), } } }