tenant.openagents/omega

1use anthropic::{
2    ANTHROPIC_API_URL, Event, Message, Request as AnthropicRequest, RequestContent,
3    Response as AnthropicResponse, ResponseContent, Role, non_streaming_completion,
4    stream_completion,
5};
6use anyhow::Result;
7use futures::StreamExt as _;
8use http_client::HttpClient;
9use indoc::indoc;
10use reqwest_client::ReqwestClient;
11use sqlez::bindable::Bind;
12use sqlez::bindable::StaticColumnCount;
13use sqlez_macros::sql;
14use std::collections::HashSet;
15use std::hash::Hash;
16use std::hash::Hasher;
17use std::path::Path;
18use std::sync::{Arc, Mutex};
19
20pub struct PlainLlmClient {
21    pub http_client: Arc<dyn HttpClient>,
22    pub api_key: String,
23}
24
25impl PlainLlmClient {
26    pub fn new() -> Result<Self> {
27        let http_client: Arc<dyn http_client::HttpClient> = Arc::new(ReqwestClient::new());
28        let api_key = std::env::var("ANTHROPIC_API_KEY")
29            .map_err(|_| anyhow::anyhow!("ANTHROPIC_API_KEY environment variable not set"))?;
30        Ok(Self {
31            http_client,
32            api_key,
33        })
34    }
35
36    pub async fn generate(
37        &self,
38        model: &str,
39        max_tokens: u64,
40        messages: Vec<Message>,
41    ) -> Result<AnthropicResponse> {
42        let request = AnthropicRequest {
43            model: model.to_string(),
44            max_tokens,
45            messages,
46            tools: Vec::new(),
47            thinking: None,
48            tool_choice: None,
49            system: None,
50            cache_control: None,
51            context_management: None,
52            metadata: None,
53            output_config: None,
54            stop_sequences: Vec::new(),
55            speed: None,
56            temperature: None,
57            top_k: None,
58            top_p: None,
59        };
60
61        let response = non_streaming_completion(
62            self.http_client.as_ref(),
63            ANTHROPIC_API_URL,
64            &self.api_key,
65            request,
66            None,
67            &http_client::CustomHeaders::default(),
68        )
69        .await
70        .map_err(|e| anyhow::anyhow!("{:?}", e))?;
71
72        Ok(response)
73    }
74
75    pub async fn generate_streaming<F>(
76        &self,
77        model: &str,
78        max_tokens: u64,
79        messages: Vec<Message>,
80        mut on_progress: F,
81    ) -> Result<AnthropicResponse>
82    where
83        F: FnMut(usize, &str),
84    {
85        let request = AnthropicRequest {
86            model: model.to_string(),
87            max_tokens,
88            messages,
89            tools: Vec::new(),
90            thinking: None,
91            tool_choice: None,
92            system: None,
93            cache_control: None,
94            context_management: None,
95            metadata: None,
96            output_config: None,
97            stop_sequences: Vec::new(),
98            speed: None,
99            temperature: None,
100            top_k: None,
101            top_p: None,
102        };
103
104        let mut stream = stream_completion(
105            self.http_client.as_ref(),
106            ANTHROPIC_API_URL,
107            &self.api_key,
108            request,
109            None,
110            &http_client::CustomHeaders::default(),
111        )
112        .await
113        .map_err(|e| anyhow::anyhow!("{:?}", e))?;
114
115        let mut response: Option<AnthropicResponse> = None;
116        let mut text_content = String::new();
117
118        while let Some(event_result) = stream.next().await {
119            let event = event_result.map_err(|e| anyhow::anyhow!("{:?}", e))?;
120
121            match event {
122                Event::MessageStart { message } => {
123                    response = Some(message);
124                }
125                Event::ContentBlockDelta { delta, .. } => {
126                    if let anthropic::ContentDelta::TextDelta { text } = delta {
127                        text_content.push_str(&text);
128                        on_progress(text_content.len(), &text_content);
129                    }
130                }
131                _ => {}
132            }
133        }
134
135        let mut response = response.ok_or_else(|| anyhow::anyhow!("No response received"))?;
136
137        if response.content.is_empty() && !text_content.is_empty() {
138            response
139                .content
140                .push(ResponseContent::Text { text: text_content });
141        }
142
143        Ok(response)
144    }
145}
146
147pub struct BatchingLlmClient {
148    connection: Mutex<sqlez::connection::Connection>,
149    http_client: Arc<dyn HttpClient>,
150    api_key: String,
151}
152
153struct CacheRow {
154    request_hash: String,
155    request: Option<String>,
156    response: Option<String>,
157    batch_id: Option<String>,
158}
159
160impl StaticColumnCount for CacheRow {
161    fn column_count() -> usize {
162        4
163    }
164}
165
166impl Bind for CacheRow {
167    fn bind(&self, statement: &sqlez::statement::Statement, start_index: i32) -> Result<i32> {
168        let next_index = statement.bind(&self.request_hash, start_index)?;
169        let next_index = statement.bind(&self.request, next_index)?;
170        let next_index = statement.bind(&self.response, next_index)?;
171        let next_index = statement.bind(&self.batch_id, next_index)?;
172        Ok(next_index)
173    }
174}
175
176#[derive(serde::Serialize, serde::Deserialize)]
177struct SerializableRequest {
178    model: String,
179    max_tokens: u64,
180    messages: Vec<SerializableMessage>,
181}
182
183#[derive(serde::Serialize, serde::Deserialize)]
184struct SerializableMessage {
185    role: String,
186    content: String,
187}
188
189impl BatchingLlmClient {
190    fn new(cache_path: &Path) -> Result<Self> {
191        let http_client: Arc<dyn http_client::HttpClient> = Arc::new(ReqwestClient::new());
192        let api_key = std::env::var("ANTHROPIC_API_KEY")
193            .map_err(|_| anyhow::anyhow!("ANTHROPIC_API_KEY environment variable not set"))?;
194
195        let connection = sqlez::connection::Connection::open_file(&cache_path.to_str().unwrap());
196        let mut statement = sqlez::statement::Statement::prepare(
197            &connection,
198            indoc! {"
199                CREATE TABLE IF NOT EXISTS cache (
200                    request_hash TEXT PRIMARY KEY,
201                    request TEXT,
202                    response TEXT,
203                    batch_id TEXT
204                );
205                "},
206        )?;
207        statement.exec()?;
208        drop(statement);
209
210        Ok(Self {
211            connection: Mutex::new(connection),
212            http_client,
213            api_key,
214        })
215    }
216
217    pub fn lookup(
218        &self,
219        model: &str,
220        max_tokens: u64,
221        messages: &[Message],
222        seed: Option<usize>,
223    ) -> Result<Option<AnthropicResponse>> {
224        let request_hash_str = Self::request_hash(model, max_tokens, messages, seed);
225        let connection = self.connection.lock().unwrap();
226        let response: Vec<String> = connection.select_bound(
227            &sql!(SELECT response FROM cache WHERE request_hash = ?1 AND response IS NOT NULL;),
228        )?(request_hash_str.as_str())?;
229        Ok(response
230            .into_iter()
231            .next()
232            .and_then(|text| serde_json::from_str(&text).ok()))
233    }
234
235    pub fn mark_for_batch(
236        &self,
237        model: &str,
238        max_tokens: u64,
239        messages: &[Message],
240        seed: Option<usize>,
241    ) -> Result<()> {
242        let request_hash = Self::request_hash(model, max_tokens, messages, seed);
243
244        let serializable_messages: Vec<SerializableMessage> = messages
245            .iter()
246            .map(|msg| SerializableMessage {
247                role: match msg.role {
248                    Role::User => "user".to_string(),
249                    Role::Assistant => "assistant".to_string(),
250                },
251                content: message_content_to_string(&msg.content),
252            })
253            .collect();
254
255        let serializable_request = SerializableRequest {
256            model: model.to_string(),
257            max_tokens,
258            messages: serializable_messages,
259        };
260
261        let request = Some(serde_json::to_string(&serializable_request)?);
262        let cache_row = CacheRow {
263            request_hash,
264            request,
265            response: None,
266            batch_id: None,
267        };
268        let connection = self.connection.lock().unwrap();
269        connection.exec_bound::<CacheRow>(sql!(
270            INSERT OR IGNORE INTO cache(request_hash, request, response, batch_id) VALUES (?, ?, ?, ?)))?(
271            cache_row,
272        )
273    }
274
275    async fn generate(
276        &self,
277        model: &str,
278        max_tokens: u64,
279        messages: Vec<Message>,
280        seed: Option<usize>,
281        cache_only: bool,
282    ) -> Result<Option<AnthropicResponse>> {
283        let response = self.lookup(model, max_tokens, &messages, seed)?;
284        if let Some(response) = response {
285            return Ok(Some(response));
286        }
287
288        if !cache_only {
289            self.mark_for_batch(model, max_tokens, &messages, seed)?;
290        }
291
292        Ok(None)
293    }
294
295    /// Uploads pending requests as batches (chunked to 16k each); downloads finished batches if any.
296    async fn sync_batches(&self) -> Result<()> {
297        let _batch_ids = self.upload_pending_requests().await?;
298        self.download_finished_batches().await
299    }
300
301    pub fn pending_batch_count(&self) -> Result<usize> {
302        let connection = self.connection.lock().unwrap();
303        let counts: Vec<i32> = connection.select(
304            sql!(SELECT COUNT(*) FROM cache WHERE batch_id IS NOT NULL AND response IS NULL),
305        )?()?;
306        Ok(counts.into_iter().next().unwrap_or(0) as usize)
307    }
308
309    /// Import batch results from external batch IDs (useful for recovering after database loss)
310    pub async fn import_batches(&self, batch_ids: &[String]) -> Result<()> {
311        for batch_id in batch_ids {
312            log::info!("Importing batch {}", batch_id);
313
314            let batch_status = anthropic::batches::retrieve_batch(
315                self.http_client.as_ref(),
316                ANTHROPIC_API_URL,
317                &self.api_key,
318                batch_id,
319            )
320            .await
321            .map_err(|e| anyhow::anyhow!("Failed to retrieve batch {}: {:?}", batch_id, e))?;
322
323            log::info!(
324                "Batch {} status: {}",
325                batch_id,
326                batch_status.processing_status
327            );
328
329            if batch_status.processing_status != "ended" {
330                log::warn!(
331                    "Batch {} is not finished (status: {}), skipping",
332                    batch_id,
333                    batch_status.processing_status
334                );
335                continue;
336            }
337
338            let results = anthropic::batches::retrieve_batch_results(
339                self.http_client.as_ref(),
340                ANTHROPIC_API_URL,
341                &self.api_key,
342                batch_id,
343            )
344            .await
345            .map_err(|e| {
346                anyhow::anyhow!("Failed to retrieve batch results for {}: {:?}", batch_id, e)
347            })?;
348
349            let mut updates: Vec<(String, String, String)> = Vec::new();
350            let mut success_count = 0;
351            let mut error_count = 0;
352
353            for result in results {
354                let request_hash = result
355                    .custom_id
356                    .strip_prefix("req_hash_")
357                    .unwrap_or(&result.custom_id)
358                    .to_string();
359
360                match result.result {
361                    anthropic::batches::BatchResult::Succeeded { message } => {
362                        let response_json = serde_json::to_string(&message)?;
363                        updates.push((request_hash, response_json, batch_id.clone()));
364                        success_count += 1;
365                    }
366                    anthropic::batches::BatchResult::Errored { error } => {
367                        log::error!(
368                            "Batch request {} failed: {}: {}",
369                            request_hash,
370                            error.error.error_type,
371                            error.error.message
372                        );
373                        let error_json = serde_json::json!({
374                            "error": {
375                                "type": error.error.error_type,
376                                "message": error.error.message
377                            }
378                        })
379                        .to_string();
380                        updates.push((request_hash, error_json, batch_id.clone()));
381                        error_count += 1;
382                    }
383                    anthropic::batches::BatchResult::Canceled => {
384                        log::warn!("Batch request {} was canceled", request_hash);
385                        error_count += 1;
386                    }
387                    anthropic::batches::BatchResult::Expired => {
388                        log::warn!("Batch request {} expired", request_hash);
389                        error_count += 1;
390                    }
391                }
392            }
393
394            let connection = self.connection.lock().unwrap();
395            connection.with_savepoint("batch_import", || {
396                // Use INSERT OR REPLACE to handle both new entries and updating existing ones
397                let q = sql!(
398                    INSERT OR REPLACE INTO cache(request_hash, request, response, batch_id)
399                    VALUES (?, (SELECT request FROM cache WHERE request_hash = ?), ?, ?)
400                );
401                let mut exec = connection.exec_bound::<(&str, &str, &str, &str)>(q)?;
402                for (request_hash, response_json, batch_id) in &updates {
403                    exec((
404                        request_hash.as_str(),
405                        request_hash.as_str(),
406                        response_json.as_str(),
407                        batch_id.as_str(),
408                    ))?;
409                }
410                Ok(())
411            })?;
412
413            log::info!(
414                "Imported batch {}: {} successful, {} errors",
415                batch_id,
416                success_count,
417                error_count
418            );
419        }
420
421        Ok(())
422    }
423
424    async fn download_finished_batches(&self) -> Result<()> {
425        let batch_ids: Vec<String> = {
426            let connection = self.connection.lock().unwrap();
427            let q = sql!(SELECT DISTINCT batch_id FROM cache WHERE batch_id IS NOT NULL AND response IS NULL);
428            connection.select(q)?()?
429        };
430
431        for batch_id in &batch_ids {
432            let batch_status = anthropic::batches::retrieve_batch(
433                self.http_client.as_ref(),
434                ANTHROPIC_API_URL,
435                &self.api_key,
436                &batch_id,
437            )
438            .await
439            .map_err(|e| anyhow::anyhow!("{:?}", e))?;
440
441            log::info!(
442                "Batch {} status: {}",
443                batch_id,
444                batch_status.processing_status
445            );
446
447            if batch_status.processing_status == "ended" {
448                let results = anthropic::batches::retrieve_batch_results(
449                    self.http_client.as_ref(),
450                    ANTHROPIC_API_URL,
451                    &self.api_key,
452                    &batch_id,
453                )
454                .await
455                .map_err(|e| anyhow::anyhow!("{:?}", e))?;
456
457                let mut updates: Vec<(String, String)> = Vec::new();
458                let mut success_count = 0;
459                for result in results {
460                    let request_hash = result
461                        .custom_id
462                        .strip_prefix("req_hash_")
463                        .unwrap_or(&result.custom_id)
464                        .to_string();
465
466                    match result.result {
467                        anthropic::batches::BatchResult::Succeeded { message } => {
468                            let response_json = serde_json::to_string(&message)?;
469                            updates.push((response_json, request_hash));
470                            success_count += 1;
471                        }
472                        anthropic::batches::BatchResult::Errored { error } => {
473                            log::error!(
474                                "Batch request {} failed: {}: {}",
475                                request_hash,
476                                error.error.error_type,
477                                error.error.message
478                            );
479                            let error_json = serde_json::json!({
480                                "error": {
481                                    "type": error.error.error_type,
482                                    "message": error.error.message
483                                }
484                            })
485                            .to_string();
486                            updates.push((error_json, request_hash));
487                        }
488                        anthropic::batches::BatchResult::Canceled => {
489                            log::warn!("Batch request {} was canceled", request_hash);
490                            let error_json = serde_json::json!({
491                                "error": {
492                                    "type": "canceled",
493                                    "message": "Batch request was canceled"
494                                }
495                            })
496                            .to_string();
497                            updates.push((error_json, request_hash));
498                        }
499                        anthropic::batches::BatchResult::Expired => {
500                            log::warn!("Batch request {} expired", request_hash);
501                            let error_json = serde_json::json!({
502                                "error": {
503                                    "type": "expired",
504                                    "message": "Batch request expired"
505                                }
506                            })
507                            .to_string();
508                            updates.push((error_json, request_hash));
509                        }
510                    }
511                }
512
513                let connection = self.connection.lock().unwrap();
514                connection.with_savepoint("batch_download", || {
515                    let q = sql!(UPDATE cache SET response = ? WHERE request_hash = ?);
516                    let mut exec = connection.exec_bound::<(&str, &str)>(q)?;
517                    for (response_json, request_hash) in &updates {
518                        exec((response_json.as_str(), request_hash.as_str()))?;
519                    }
520                    Ok(())
521                })?;
522                log::info!("Downloaded {} successful requests", success_count);
523            }
524        }
525
526        Ok(())
527    }
528
529    async fn upload_pending_requests(&self) -> Result<Vec<String>> {
530        const BATCH_CHUNK_SIZE: i32 = 16_000;
531        const MAX_BATCH_SIZE_BYTES: usize = 100 * 1024 * 1024;
532
533        let mut all_batch_ids = Vec::new();
534        let mut total_uploaded = 0;
535
536        let mut current_batch_rows = Vec::new();
537        let mut current_batch_size = 0usize;
538        let mut pending_hashes: HashSet<String> = HashSet::new();
539        loop {
540            let rows: Vec<(String, String)> = {
541                let connection = self.connection.lock().unwrap();
542                let q = sql!(
543                    SELECT request_hash, request FROM cache
544                    WHERE batch_id IS NULL AND response IS NULL
545                    LIMIT ?
546                );
547                connection.select_bound(q)?(BATCH_CHUNK_SIZE)?
548            };
549
550            if rows.is_empty() {
551                break;
552            }
553
554            // Split rows into sub-batches based on size
555            let mut batches_to_upload = Vec::new();
556            let mut new_rows_added = 0;
557
558            for row in rows {
559                let (hash, request_str) = row;
560
561                // Skip rows already added to current_batch_rows but not yet uploaded
562                if pending_hashes.contains(&hash) {
563                    continue;
564                }
565                let serializable_request: SerializableRequest = serde_json::from_str(&request_str)?;
566
567                let messages: Vec<Message> = serializable_request
568                    .messages
569                    .into_iter()
570                    .map(|msg| Message {
571                        role: match msg.role.as_str() {
572                            "user" => Role::User,
573                            "assistant" => Role::Assistant,
574                            _ => Role::User,
575                        },
576                        content: vec![RequestContent::Text {
577                            text: msg.content,
578                            cache_control: None,
579                        }],
580                    })
581                    .collect();
582
583                let params = AnthropicRequest {
584                    model: serializable_request.model,
585                    max_tokens: serializable_request.max_tokens,
586                    messages,
587                    tools: Vec::new(),
588                    thinking: None,
589                    tool_choice: None,
590                    system: None,
591                    cache_control: None,
592                    context_management: None,
593                    metadata: None,
594                    output_config: None,
595                    stop_sequences: Vec::new(),
596                    temperature: None,
597                    top_k: None,
598                    top_p: None,
599                    speed: None,
600                };
601
602                let custom_id = format!("req_hash_{}", hash);
603                let batch_request = anthropic::batches::BatchRequest { custom_id, params };
604
605                // Estimate the serialized size of this request
606                let estimated_size = serde_json::to_string(&batch_request)?.len();
607
608                // If adding this request would exceed the limit, start a new batch
609                if !current_batch_rows.is_empty()
610                    && current_batch_size + estimated_size > MAX_BATCH_SIZE_BYTES
611                {
612                    batches_to_upload.push((current_batch_rows, current_batch_size));
613                    current_batch_rows = Vec::new();
614                    current_batch_size = 0;
615                }
616
617                pending_hashes.insert(hash.clone());
618                current_batch_rows.push((hash, batch_request));
619                current_batch_size += estimated_size;
620                new_rows_added += 1;
621            }
622
623            // If no new rows were added this iteration, all pending requests are already
624            // in current_batch_rows, so we should break to avoid an infinite loop
625            if new_rows_added == 0 {
626                break;
627            }
628
629            // Only upload full batches, keep the partial batch for the next iteration
630            // Upload each sub-batch
631            for (batch_rows, batch_size) in batches_to_upload {
632                let request_hashes: Vec<String> =
633                    batch_rows.iter().map(|(hash, _)| hash.clone()).collect();
634
635                // Remove uploaded hashes from pending set
636                for hash in &request_hashes {
637                    pending_hashes.remove(hash);
638                }
639                let batch_requests: Vec<anthropic::batches::BatchRequest> =
640                    batch_rows.into_iter().map(|(_, req)| req).collect();
641
642                let batch_len = batch_requests.len();
643                log::info!(
644                    "Uploading batch with {} requests (~{:.2} MB)",
645                    batch_len,
646                    batch_size as f64 / (1024.0 * 1024.0)
647                );
648
649                let batch = anthropic::batches::create_batch(
650                    self.http_client.as_ref(),
651                    ANTHROPIC_API_URL,
652                    &self.api_key,
653                    anthropic::batches::CreateBatchRequest {
654                        requests: batch_requests,
655                    },
656                )
657                .await
658                .map_err(|e| anyhow::anyhow!("{:?}", e))?;
659
660                {
661                    let connection = self.connection.lock().unwrap();
662                    connection.with_savepoint("batch_upload", || {
663                        let q = sql!(UPDATE cache SET batch_id = ? WHERE request_hash = ?);
664                        let mut exec = connection.exec_bound::<(&str, &str)>(q)?;
665                        for hash in &request_hashes {
666                            exec((batch.id.as_str(), hash.as_str()))?;
667                        }
668                        Ok(())
669                    })?;
670                }
671
672                total_uploaded += batch_len;
673                log::info!(
674                    "Uploaded batch {} with {} requests ({} total)",
675                    batch.id,
676                    batch_len,
677                    total_uploaded
678                );
679
680                all_batch_ids.push(batch.id);
681            }
682        }
683
684        // Upload any remaining partial batch at the end
685        if !current_batch_rows.is_empty() {
686            let request_hashes: Vec<String> = current_batch_rows
687                .iter()
688                .map(|(hash, _)| hash.clone())
689                .collect();
690            let batch_requests: Vec<anthropic::batches::BatchRequest> =
691                current_batch_rows.into_iter().map(|(_, req)| req).collect();
692
693            let batch_len = batch_requests.len();
694            log::info!(
695                "Uploading final batch with {} requests (~{:.2} MB)",
696                batch_len,
697                current_batch_size as f64 / (1024.0 * 1024.0)
698            );
699
700            let batch = anthropic::batches::create_batch(
701                self.http_client.as_ref(),
702                ANTHROPIC_API_URL,
703                &self.api_key,
704                anthropic::batches::CreateBatchRequest {
705                    requests: batch_requests,
706                },
707            )
708            .await
709            .map_err(|e| anyhow::anyhow!("{:?}", e))?;
710
711            {
712                let connection = self.connection.lock().unwrap();
713                connection.with_savepoint("batch_upload", || {
714                    let q = sql!(UPDATE cache SET batch_id = ? WHERE request_hash = ?);
715                    let mut exec = connection.exec_bound::<(&str, &str)>(q)?;
716                    for hash in &request_hashes {
717                        exec((batch.id.as_str(), hash.as_str()))?;
718                    }
719                    Ok(())
720                })?;
721            }
722
723            total_uploaded += batch_len;
724            log::info!(
725                "Uploaded batch {} with {} requests ({} total)",
726                batch.id,
727                batch_len,
728                total_uploaded
729            );
730
731            all_batch_ids.push(batch.id);
732        }
733
734        if !all_batch_ids.is_empty() {
735            log::info!(
736                "Finished uploading {} batches with {} total requests",
737                all_batch_ids.len(),
738                total_uploaded
739            );
740        }
741
742        Ok(all_batch_ids)
743    }
744
745    fn request_hash(
746        model: &str,
747        max_tokens: u64,
748        messages: &[Message],
749        seed: Option<usize>,
750    ) -> String {
751        let mut hasher = std::hash::DefaultHasher::new();
752        model.hash(&mut hasher);
753        max_tokens.hash(&mut hasher);
754        for msg in messages {
755            message_content_to_string(&msg.content).hash(&mut hasher);
756        }
757        if let Some(seed) = seed {
758            seed.hash(&mut hasher);
759        }
760        let request_hash = hasher.finish();
761        format!("{request_hash:016x}")
762    }
763}
764
765fn message_content_to_string(content: &[RequestContent]) -> String {
766    content
767        .iter()
768        .filter_map(|c| match c {
769            RequestContent::Text { text, .. } => Some(text.clone()),
770            _ => None,
771        })
772        .collect::<Vec<String>>()
773        .join("\n")
774}
775
776pub enum AnthropicClient {
777    // No batching
778    Plain(PlainLlmClient),
779    Batch(BatchingLlmClient),
780    Dummy,
781}
782
783impl AnthropicClient {
784    pub fn plain() -> Result<Self> {
785        Ok(Self::Plain(PlainLlmClient::new()?))
786    }
787
788    pub fn batch(cache_path: &Path) -> Result<Self> {
789        Ok(Self::Batch(BatchingLlmClient::new(cache_path)?))
790    }
791
792    #[allow(dead_code)]
793    pub fn dummy() -> Self {
794        Self::Dummy
795    }
796
797    pub async fn generate(
798        &self,
799        model: &str,
800        max_tokens: u64,
801        messages: Vec<Message>,
802        seed: Option<usize>,
803        cache_only: bool,
804    ) -> Result<Option<AnthropicResponse>> {
805        match self {
806            AnthropicClient::Plain(plain_llm_client) => plain_llm_client
807                .generate(model, max_tokens, messages)
808                .await
809                .map(Some),
810            AnthropicClient::Batch(batching_llm_client) => {
811                batching_llm_client
812                    .generate(model, max_tokens, messages, seed, cache_only)
813                    .await
814            }
815            AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
816        }
817    }
818
819    #[allow(dead_code)]
820    pub async fn generate_streaming<F>(
821        &self,
822        model: &str,
823        max_tokens: u64,
824        messages: Vec<Message>,
825        on_progress: F,
826    ) -> Result<Option<AnthropicResponse>>
827    where
828        F: FnMut(usize, &str),
829    {
830        match self {
831            AnthropicClient::Plain(plain_llm_client) => plain_llm_client
832                .generate_streaming(model, max_tokens, messages, on_progress)
833                .await
834                .map(Some),
835            AnthropicClient::Batch(_) => {
836                anyhow::bail!("Streaming not supported with batching client")
837            }
838            AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
839        }
840    }
841
842    pub async fn sync_batches(&self) -> Result<()> {
843        match self {
844            AnthropicClient::Plain(_) => Ok(()),
845            AnthropicClient::Batch(batching_llm_client) => batching_llm_client.sync_batches().await,
846            AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
847        }
848    }
849
850    pub fn pending_batch_count(&self) -> Result<usize> {
851        match self {
852            AnthropicClient::Plain(_) => Ok(0),
853            AnthropicClient::Batch(batching_llm_client) => {
854                batching_llm_client.pending_batch_count()
855            }
856            AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
857        }
858    }
859
860    pub async fn import_batches(&self, batch_ids: &[String]) -> Result<()> {
861        match self {
862            AnthropicClient::Plain(_) => {
863                anyhow::bail!("Import batches is only supported with batching client")
864            }
865            AnthropicClient::Batch(batching_llm_client) => {
866                batching_llm_client.import_batches(batch_ids).await
867            }
868            AnthropicClient::Dummy => panic!("Dummy LLM client is not expected to be used"),
869        }
870    }
871}
872