tenant.openagents/omega

1use anyhow::Result;
2use http_client::HttpClient;
3use indoc::indoc;
4use open_ai::{
5    MessageContent, OPEN_AI_API_URL, Request as OpenAiRequest, RequestMessage,
6    Response as OpenAiResponse, batches, non_streaming_completion,
7};
8use reqwest_client::ReqwestClient;
9use sqlez::bindable::Bind;
10use sqlez::bindable::StaticColumnCount;
11use sqlez_macros::sql;
12use std::hash::Hash;
13use std::hash::Hasher;
14use std::path::Path;
15use std::sync::{Arc, Mutex};
16
17pub struct PlainOpenAiClient {
18    pub http_client: Arc<dyn HttpClient>,
19    pub api_key: String,
20}
21
22impl PlainOpenAiClient {
23    pub fn new() -> Result<Self> {
24        let http_client: Arc<dyn http_client::HttpClient> = Arc::new(ReqwestClient::new());
25        let api_key = std::env::var("OPENAI_API_KEY")
26            .map_err(|_| anyhow::anyhow!("OPENAI_API_KEY environment variable not set"))?;
27        Ok(Self {
28            http_client,
29            api_key,
30        })
31    }
32
33    pub async fn generate(
34        &self,
35        model: &str,
36        max_tokens: u64,
37        messages: Vec<RequestMessage>,
38    ) -> Result<OpenAiResponse> {
39        let request = OpenAiRequest {
40            model: model.to_string(),
41            messages,
42            stream: false,
43            stream_options: None,
44            max_completion_tokens: Some(max_tokens),
45            max_tokens: None,
46            stop: Vec::new(),
47            temperature: None,
48            tool_choice: None,
49            parallel_tool_calls: None,
50            service_tier: None,
51            tools: Vec::new(),
52            prompt_cache_key: None,
53            reasoning_effort: None,
54        };
55
56        let response = non_streaming_completion(
57            self.http_client.as_ref(),
58            OPEN_AI_API_URL,
59            &self.api_key,
60            request,
61        )
62        .await
63        .map_err(|e| anyhow::anyhow!("{:?}", e))?;
64
65        Ok(response)
66    }
67}
68
69pub struct BatchingOpenAiClient {
70    connection: Mutex<sqlez::connection::Connection>,
71    http_client: Arc<dyn HttpClient>,
72    api_key: String,
73}
74
75struct CacheRow {
76    request_hash: String,
77    request: Option<String>,
78    response: Option<String>,
79    batch_id: Option<String>,
80}
81
82impl StaticColumnCount for CacheRow {
83    fn column_count() -> usize {
84        4
85    }
86}
87
88impl Bind for CacheRow {
89    fn bind(&self, statement: &sqlez::statement::Statement, start_index: i32) -> Result<i32> {
90        let next_index = statement.bind(&self.request_hash, start_index)?;
91        let next_index = statement.bind(&self.request, next_index)?;
92        let next_index = statement.bind(&self.response, next_index)?;
93        let next_index = statement.bind(&self.batch_id, next_index)?;
94        Ok(next_index)
95    }
96}
97
98#[derive(serde::Serialize, serde::Deserialize)]
99struct SerializableRequest {
100    model: String,
101    max_tokens: u64,
102    messages: Vec<SerializableMessage>,
103}
104
105#[derive(serde::Serialize, serde::Deserialize)]
106struct SerializableMessage {
107    role: String,
108    content: String,
109}
110
111impl BatchingOpenAiClient {
112    fn new(cache_path: &Path) -> Result<Self> {
113        let http_client: Arc<dyn http_client::HttpClient> = Arc::new(ReqwestClient::new());
114        let api_key = std::env::var("OPENAI_API_KEY")
115            .map_err(|_| anyhow::anyhow!("OPENAI_API_KEY environment variable not set"))?;
116
117        let connection = sqlez::connection::Connection::open_file(cache_path.to_str().unwrap());
118        let mut statement = sqlez::statement::Statement::prepare(
119            &connection,
120            indoc! {"
121                CREATE TABLE IF NOT EXISTS openai_cache (
122                    request_hash TEXT PRIMARY KEY,
123                    request TEXT,
124                    response TEXT,
125                    batch_id TEXT
126                );
127                "},
128        )?;
129        statement.exec()?;
130        drop(statement);
131
132        Ok(Self {
133            connection: Mutex::new(connection),
134            http_client,
135            api_key,
136        })
137    }
138
139    pub fn lookup(
140        &self,
141        model: &str,
142        max_tokens: u64,
143        messages: &[RequestMessage],
144        seed: Option<usize>,
145    ) -> Result<Option<OpenAiResponse>> {
146        let request_hash_str = Self::request_hash(model, max_tokens, messages, seed);
147        let connection = self.connection.lock().unwrap();
148        let response: Vec<String> = connection.select_bound(
149            &sql!(SELECT response FROM openai_cache WHERE request_hash = ?1 AND response IS NOT NULL;),
150        )?(request_hash_str.as_str())?;
151        Ok(response
152            .into_iter()
153            .next()
154            .and_then(|text| serde_json::from_str(&text).ok()))
155    }
156
157    pub fn mark_for_batch(
158        &self,
159        model: &str,
160        max_tokens: u64,
161        messages: &[RequestMessage],
162        seed: Option<usize>,
163    ) -> Result<()> {
164        let request_hash = Self::request_hash(model, max_tokens, messages, seed);
165
166        let serializable_messages: Vec<SerializableMessage> = messages
167            .iter()
168            .map(|msg| SerializableMessage {
169                role: message_role_to_string(msg),
170                content: message_content_to_string(msg),
171            })
172            .collect();
173
174        let serializable_request = SerializableRequest {
175            model: model.to_string(),
176            max_tokens,
177            messages: serializable_messages,
178        };
179
180        let request = Some(serde_json::to_string(&serializable_request)?);
181        let cache_row = CacheRow {
182            request_hash,
183            request,
184            response: None,
185            batch_id: None,
186        };
187        let connection = self.connection.lock().unwrap();
188        connection.exec_bound::<CacheRow>(sql!(
189            INSERT OR IGNORE INTO openai_cache(request_hash, request, response, batch_id) VALUES (?, ?, ?, ?)))?(
190            cache_row,
191        )
192    }
193
194    async fn generate(
195        &self,
196        model: &str,
197        max_tokens: u64,
198        messages: Vec<RequestMessage>,
199        seed: Option<usize>,
200        cache_only: bool,
201    ) -> Result<Option<OpenAiResponse>> {
202        let response = self.lookup(model, max_tokens, &messages, seed)?;
203        if let Some(response) = response {
204            return Ok(Some(response));
205        }
206
207        if !cache_only {
208            self.mark_for_batch(model, max_tokens, &messages, seed)?;
209        }
210
211        Ok(None)
212    }
213
214    async fn sync_batches(&self) -> Result<()> {
215        let _batch_ids = self.upload_pending_requests().await?;
216        self.download_finished_batches().await
217    }
218
219    pub fn pending_batch_count(&self) -> Result<usize> {
220        let connection = self.connection.lock().unwrap();
221        let counts: Vec<i32> = connection.select(
222            sql!(SELECT COUNT(*) FROM openai_cache WHERE batch_id IS NOT NULL AND response IS NULL),
223        )?()?;
224        Ok(counts.into_iter().next().unwrap_or(0) as usize)
225    }
226
227    pub async fn import_batches(&self, batch_ids: &[String]) -> Result<()> {
228        for batch_id in batch_ids {
229            log::info!("Importing OpenAI batch {}", batch_id);
230
231            let batch_status = batches::retrieve_batch(
232                self.http_client.as_ref(),
233                OPEN_AI_API_URL,
234                &self.api_key,
235                batch_id,
236            )
237            .await
238            .map_err(|e| anyhow::anyhow!("Failed to retrieve batch {}: {:?}", batch_id, e))?;
239
240            log::info!("Batch {} status: {}", batch_id, batch_status.status);
241
242            if batch_status.status != "completed" {
243                log::warn!(
244                    "Batch {} is not completed (status: {}), skipping",
245                    batch_id,
246                    batch_status.status
247                );
248                continue;
249            }
250
251            let output_file_id = batch_status.output_file_id.ok_or_else(|| {
252                anyhow::anyhow!("Batch {} completed but has no output file", batch_id)
253            })?;
254
255            let results_content = batches::download_file(
256                self.http_client.as_ref(),
257                OPEN_AI_API_URL,
258                &self.api_key,
259                &output_file_id,
260            )
261            .await
262            .map_err(|e| {
263                anyhow::anyhow!("Failed to download batch results for {}: {:?}", batch_id, e)
264            })?;
265
266            let results = batches::parse_batch_output(&results_content)
267                .map_err(|e| anyhow::anyhow!("Failed to parse batch output: {:?}", e))?;
268
269            let mut updates: Vec<(String, String, String)> = Vec::new();
270            let mut success_count = 0;
271            let mut error_count = 0;
272
273            for result in results {
274                let request_hash = result
275                    .custom_id
276                    .strip_prefix("req_hash_")
277                    .unwrap_or(&result.custom_id)
278                    .to_string();
279
280                if let Some(response_body) = result.response {
281                    if response_body.status_code == 200 {
282                        let response_json = serde_json::to_string(&response_body.body)?;
283                        updates.push((request_hash, response_json, batch_id.clone()));
284                        success_count += 1;
285                    } else {
286                        log::error!(
287                            "Batch request {} failed with status {}",
288                            request_hash,
289                            response_body.status_code
290                        );
291                        let error_json = serde_json::json!({
292                            "error": {
293                                "type": "http_error",
294                                "status_code": response_body.status_code
295                            }
296                        })
297                        .to_string();
298                        updates.push((request_hash, error_json, batch_id.clone()));
299                        error_count += 1;
300                    }
301                } else if let Some(error) = result.error {
302                    log::error!(
303                        "Batch request {} failed: {}: {}",
304                        request_hash,
305                        error.code,
306                        error.message
307                    );
308                    let error_json = serde_json::json!({
309                        "error": {
310                            "type": error.code,
311                            "message": error.message
312                        }
313                    })
314                    .to_string();
315                    updates.push((request_hash, error_json, batch_id.clone()));
316                    error_count += 1;
317                }
318            }
319
320            let connection = self.connection.lock().unwrap();
321            connection.with_savepoint("batch_import", || {
322                let q = sql!(
323                    INSERT OR REPLACE INTO openai_cache(request_hash, request, response, batch_id)
324                    VALUES (?, (SELECT request FROM openai_cache WHERE request_hash = ?), ?, ?)
325                );
326                let mut exec = connection.exec_bound::<(&str, &str, &str, &str)>(q)?;
327                for (request_hash, response_json, batch_id) in &updates {
328                    exec((
329                        request_hash.as_str(),
330                        request_hash.as_str(),
331                        response_json.as_str(),
332                        batch_id.as_str(),
333                    ))?;
334                }
335                Ok(())
336            })?;
337
338            log::info!(
339                "Imported batch {}: {} successful, {} errors",
340                batch_id,
341                success_count,
342                error_count
343            );
344        }
345
346        Ok(())
347    }
348
349    async fn download_finished_batches(&self) -> Result<()> {
350        let batch_ids: Vec<String> = {
351            let connection = self.connection.lock().unwrap();
352            let q = sql!(SELECT DISTINCT batch_id FROM openai_cache WHERE batch_id IS NOT NULL AND response IS NULL);
353            connection.select(q)?()?
354        };
355
356        for batch_id in &batch_ids {
357            let batch_status = batches::retrieve_batch(
358                self.http_client.as_ref(),
359                OPEN_AI_API_URL,
360                &self.api_key,
361                batch_id,
362            )
363            .await
364            .map_err(|e| anyhow::anyhow!("{:?}", e))?;
365
366            log::info!("Batch {} status: {}", batch_id, batch_status.status);
367
368            if batch_status.status == "completed" {
369                let output_file_id = match batch_status.output_file_id {
370                    Some(id) => id,
371                    None => {
372                        log::warn!("Batch {} completed but has no output file", batch_id);
373                        continue;
374                    }
375                };
376
377                let results_content = batches::download_file(
378                    self.http_client.as_ref(),
379                    OPEN_AI_API_URL,
380                    &self.api_key,
381                    &output_file_id,
382                )
383                .await
384                .map_err(|e| anyhow::anyhow!("{:?}", e))?;
385
386                let results = batches::parse_batch_output(&results_content)
387                    .map_err(|e| anyhow::anyhow!("Failed to parse batch output: {:?}", e))?;
388
389                let mut updates: Vec<(String, String)> = Vec::new();
390                let mut success_count = 0;
391
392                for result in results {
393                    let request_hash = result
394                        .custom_id
395                        .strip_prefix("req_hash_")
396                        .unwrap_or(&result.custom_id)
397                        .to_string();
398
399                    if let Some(response_body) = result.response {
400                        if response_body.status_code == 200 {
401                            let response_json = serde_json::to_string(&response_body.body)?;
402                            updates.push((response_json, request_hash));
403                            success_count += 1;
404                        } else {
405                            log::error!(
406                                "Batch request {} failed with status {}",
407                                request_hash,
408                                response_body.status_code
409                            );
410                            let error_json = serde_json::json!({
411                                "error": {
412                                    "type": "http_error",
413                                    "status_code": response_body.status_code
414                                }
415                            })
416                            .to_string();
417                            updates.push((error_json, request_hash));
418                        }
419                    } else if let Some(error) = result.error {
420                        log::error!(
421                            "Batch request {} failed: {}: {}",
422                            request_hash,
423                            error.code,
424                            error.message
425                        );
426                        let error_json = serde_json::json!({
427                            "error": {
428                                "type": error.code,
429                                "message": error.message
430                            }
431                        })
432                        .to_string();
433                        updates.push((error_json, request_hash));
434                    }
435                }
436
437                let connection = self.connection.lock().unwrap();
438                connection.with_savepoint("batch_download", || {
439                    let q = sql!(UPDATE openai_cache SET response = ? WHERE request_hash = ?);
440                    let mut exec = connection.exec_bound::<(&str, &str)>(q)?;
441                    for (response_json, request_hash) in &updates {
442                        exec((response_json.as_str(), request_hash.as_str()))?;
443                    }
444                    Ok(())
445                })?;
446                log::info!("Downloaded {} successful requests", success_count);
447            }
448        }
449
450        Ok(())
451    }
452
453    async fn upload_pending_requests(&self) -> Result<Vec<String>> {
454        const BATCH_CHUNK_SIZE: i32 = 16_000;
455        let mut all_batch_ids = Vec::new();
456        let mut total_uploaded = 0;
457
458        loop {
459            let rows: Vec<(String, String)> = {
460                let connection = self.connection.lock().unwrap();
461                let q = sql!(
462                    SELECT request_hash, request FROM openai_cache
463                    WHERE batch_id IS NULL AND response IS NULL
464                    LIMIT ?
465                );
466                connection.select_bound(q)?(BATCH_CHUNK_SIZE)?
467            };
468
469            if rows.is_empty() {
470                break;
471            }
472
473            let request_hashes: Vec<String> = rows.iter().map(|(hash, _)| hash.clone()).collect();
474
475            let mut jsonl_content = String::new();
476            for (hash, request_str) in &rows {
477                let serializable_request: SerializableRequest =
478                    serde_json::from_str(request_str).unwrap();
479
480                let messages: Vec<RequestMessage> = serializable_request
481                    .messages
482                    .into_iter()
483                    .map(|msg| match msg.role.as_str() {
484                        "user" => RequestMessage::User {
485                            content: MessageContent::Plain(msg.content),
486                        },
487                        "assistant" => RequestMessage::Assistant {
488                            content: Some(MessageContent::Plain(msg.content)),
489                            tool_calls: Vec::new(),
490                            reasoning_content: None,
491                        },
492                        "system" => RequestMessage::System {
493                            content: MessageContent::Plain(msg.content),
494                        },
495                        _ => RequestMessage::User {
496                            content: MessageContent::Plain(msg.content),
497                        },
498                    })
499                    .collect();
500
501                let request = OpenAiRequest {
502                    model: serializable_request.model,
503                    messages,
504                    stream: false,
505                    stream_options: None,
506                    max_completion_tokens: Some(serializable_request.max_tokens),
507                    max_tokens: None,
508                    stop: Vec::new(),
509                    temperature: None,
510                    tool_choice: None,
511                    parallel_tool_calls: None,
512                    service_tier: None,
513                    tools: Vec::new(),
514                    prompt_cache_key: None,
515                    reasoning_effort: None,
516                };
517
518                let custom_id = format!("req_hash_{}", hash);
519                let batch_item = batches::BatchRequestItem::new(custom_id, request);
520                let line = batch_item
521                    .to_jsonl_line()
522                    .map_err(|e| anyhow::anyhow!("Failed to serialize batch item: {:?}", e))?;
523                jsonl_content.push_str(&line);
524                jsonl_content.push('\n');
525            }
526
527            let filename = format!("batch_{}.jsonl", chrono::Utc::now().timestamp());
528            let file_obj = batches::upload_batch_file(
529                self.http_client.as_ref(),
530                OPEN_AI_API_URL,
531                &self.api_key,
532                &filename,
533                jsonl_content.into_bytes(),
534            )
535            .await
536            .map_err(|e| anyhow::anyhow!("Failed to upload batch file: {:?}", e))?;
537
538            let batch = batches::create_batch(
539                self.http_client.as_ref(),
540                OPEN_AI_API_URL,
541                &self.api_key,
542                batches::CreateBatchRequest::new(file_obj.id),
543            )
544            .await
545            .map_err(|e| anyhow::anyhow!("Failed to create batch: {:?}", e))?;
546
547            {
548                let connection = self.connection.lock().unwrap();
549                connection.with_savepoint("batch_upload", || {
550                    let q = sql!(UPDATE openai_cache SET batch_id = ? WHERE request_hash = ?);
551                    let mut exec = connection.exec_bound::<(&str, &str)>(q)?;
552                    for hash in &request_hashes {
553                        exec((batch.id.as_str(), hash.as_str()))?;
554                    }
555                    Ok(())
556                })?;
557            }
558
559            let batch_len = rows.len();
560            total_uploaded += batch_len;
561            log::info!(
562                "Uploaded batch {} with {} requests ({} total)",
563                batch.id,
564                batch_len,
565                total_uploaded
566            );
567
568            all_batch_ids.push(batch.id);
569        }
570
571        if !all_batch_ids.is_empty() {
572            log::info!(
573                "Finished uploading {} batches with {} total requests",
574                all_batch_ids.len(),
575                total_uploaded
576            );
577        }
578
579        Ok(all_batch_ids)
580    }
581
582    fn request_hash(
583        model: &str,
584        max_tokens: u64,
585        messages: &[RequestMessage],
586        seed: Option<usize>,
587    ) -> String {
588        let mut hasher = std::hash::DefaultHasher::new();
589        "openai".hash(&mut hasher);
590        model.hash(&mut hasher);
591        max_tokens.hash(&mut hasher);
592        for msg in messages {
593            message_content_to_string(msg).hash(&mut hasher);
594        }
595        if let Some(seed) = seed {
596            seed.hash(&mut hasher);
597        }
598        let request_hash = hasher.finish();
599        format!("{request_hash:016x}")
600    }
601}
602
603fn message_role_to_string(msg: &RequestMessage) -> String {
604    match msg {
605        RequestMessage::User { .. } => "user".to_string(),
606        RequestMessage::Assistant { .. } => "assistant".to_string(),
607        RequestMessage::System { .. } => "system".to_string(),
608        RequestMessage::Tool { .. } => "tool".to_string(),
609    }
610}
611
612fn message_content_to_string(msg: &RequestMessage) -> String {
613    match msg {
614        RequestMessage::User { content } => content_to_string(content),
615        RequestMessage::Assistant { content, .. } => {
616            content.as_ref().map(content_to_string).unwrap_or_default()
617        }
618        RequestMessage::System { content } => content_to_string(content),
619        RequestMessage::Tool { content, .. } => content_to_string(content),
620    }
621}
622
623fn content_to_string(content: &MessageContent) -> String {
624    match content {
625        MessageContent::Plain(text) => text.clone(),
626        MessageContent::Multipart(parts) => parts
627            .iter()
628            .filter_map(|part| match part {
629                open_ai::MessagePart::Text { text } => Some(text.clone()),
630                _ => None,
631            })
632            .collect::<Vec<String>>()
633            .join("\n"),
634    }
635}
636
637pub enum OpenAiClient {
638    Plain(PlainOpenAiClient),
639    Batch(BatchingOpenAiClient),
640    #[allow(dead_code)]
641    Dummy,
642}
643
644impl OpenAiClient {
645    pub fn plain() -> Result<Self> {
646        Ok(Self::Plain(PlainOpenAiClient::new()?))
647    }
648
649    pub fn batch(cache_path: &Path) -> Result<Self> {
650        Ok(Self::Batch(BatchingOpenAiClient::new(cache_path)?))
651    }
652
653    #[allow(dead_code)]
654    pub fn dummy() -> Self {
655        Self::Dummy
656    }
657
658    pub async fn generate(
659        &self,
660        model: &str,
661        max_tokens: u64,
662        messages: Vec<RequestMessage>,
663        seed: Option<usize>,
664        cache_only: bool,
665    ) -> Result<Option<OpenAiResponse>> {
666        match self {
667            OpenAiClient::Plain(plain_client) => plain_client
668                .generate(model, max_tokens, messages)
669                .await
670                .map(Some),
671            OpenAiClient::Batch(batching_client) => {
672                batching_client
673                    .generate(model, max_tokens, messages, seed, cache_only)
674                    .await
675            }
676            OpenAiClient::Dummy => panic!("Dummy OpenAI client is not expected to be used"),
677        }
678    }
679
680    pub async fn sync_batches(&self) -> Result<()> {
681        match self {
682            OpenAiClient::Plain(_) => Ok(()),
683            OpenAiClient::Batch(batching_client) => batching_client.sync_batches().await,
684            OpenAiClient::Dummy => panic!("Dummy OpenAI client is not expected to be used"),
685        }
686    }
687
688    pub fn pending_batch_count(&self) -> Result<usize> {
689        match self {
690            OpenAiClient::Plain(_) => Ok(0),
691            OpenAiClient::Batch(batching_client) => batching_client.pending_batch_count(),
692            OpenAiClient::Dummy => panic!("Dummy OpenAI client is not expected to be used"),
693        }
694    }
695
696    pub async fn import_batches(&self, batch_ids: &[String]) -> Result<()> {
697        match self {
698            OpenAiClient::Plain(_) => {
699                anyhow::bail!("Import batches is only supported with batching client")
700            }
701            OpenAiClient::Batch(batching_client) => batching_client.import_batches(batch_ids).await,
702            OpenAiClient::Dummy => panic!("Dummy OpenAI client is not expected to be used"),
703        }
704    }
705}
706