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<dyn HttpClient>, pub api_key: String,} impl PlainLlmClient { pub fn new() -> Result<Self> { let http_client: Arc<dyn http_client::HttpClient> = 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<Message>, ) -> Result<AnthropicResponse> { 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<F>( &self, model: &str, max_tokens: u64, messages: Vec<Message>, mut on_progress: F, ) -> Result<AnthropicResponse> 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<AnthropicResponse> = 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<sqlez::connection::Connection>, http_client: Arc<dyn HttpClient>, api_key: String,} struct CacheRow { request_hash: String, request: Option<String>, response: Option<String>, batch_id: Option<String>,} impl StaticColumnCount for CacheRow { fn column_count() -> usize { 4 }} impl Bind for CacheRow { fn bind(&self, statement: &sqlez::statement::Statement, start_index: i32) -> Result<i32> { 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<SerializableMessage>,} #[derive(serde::Serialize, serde::Deserialize)]struct SerializableMessage { role: String, content: String,} impl BatchingLlmClient { fn new(cache_path: &Path) -> Result<Self> { let http_client: Arc<dyn http_client::HttpClient> = 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<usize>, ) -> Result<Option<AnthropicResponse>> { let request_hash_str = Self::request_hash(model, max_tokens, messages, seed); let connection = self.connection.lock().unwrap(); let response: Vec<String> = 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<usize>, ) -> Result<()> { let request_hash = Self::request_hash(model, max_tokens, messages, seed); let serializable_messages: Vec<SerializableMessage> = 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::<CacheRow>(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<Message>, seed: Option<usize>, cache_only: bool, ) -> Result<Option<AnthropicResponse>> { 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<usize> { let connection = self.connection.lock().unwrap(); let counts: Vec<i32> = 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<String> = { 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);Lines 1–500 of 871. Load more
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<dyn HttpClient>, pub api_key: String, } impl PlainLlmClient { pub fn new() -> Result<Self> { let http_client: Arc<dyn http_client::HttpClient> = 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<Message>, ) -> Result<AnthropicResponse> { 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<F>( &self, model: &str, max_tokens: u64, messages: Vec<Message>, mut on_progress: F, ) -> Result<AnthropicResponse> 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<AnthropicResponse> = 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<sqlez::connection::Connection>, http_client: Arc<dyn HttpClient>, api_key: String, } struct CacheRow { request_hash: String, request: Option<String>, response: Option<String>, batch_id: Option<String>, } impl StaticColumnCount for CacheRow { fn column_count() -> usize { 4 } } impl Bind for CacheRow { fn bind(&self, statement: &sqlez::statement::Statement, start_index: i32) -> Result<i32> { 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<SerializableMessage>, } #[derive(serde::Serialize, serde::Deserialize)] struct SerializableMessage { role: String, content: String, } impl BatchingLlmClient { fn new(cache_path: &Path) -> Result<Self> { let http_client: Arc<dyn http_client::HttpClient> = 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<usize>, ) -> Result<Option<AnthropicResponse>> { let request_hash_str = Self::request_hash(model, max_tokens, messages, seed); let connection = self.connection.lock().unwrap(); let response: Vec<String> = 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<usize>, ) -> Result<()> { let request_hash = Self::request_hash(model, max_tokens, messages, seed); let serializable_messages: Vec<SerializableMessage> = 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::<CacheRow>(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<Message>, seed: Option<usize>, cache_only: bool, ) -> Result<Option<AnthropicResponse>> { 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<usize> { let connection = self.connection.lock().unwrap(); let counts: Vec<i32> = 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<String> = { 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);