Skip to main content

shift_proxy/routes/
openai.rs

1//! OpenAI route handler — POST /v1/chat/completions
2//!
3//! Intercepts OpenAI API requests, runs the SHIFT optimization pipeline
4//! on the payload, records stats with full per-image token savings, and
5//! forwards to the real OpenAI API.
6//!
7//! The CPU-intensive optimization runs on a blocking thread to avoid
8//! starving the tokio event loop.
9
10use crate::forward::forward_request;
11use crate::optimize::optimize_payload;
12use crate::ProxyState;
13use axum::extract::State;
14use axum::http::{HeaderMap, Uri};
15use axum::response::Response;
16
17/// POST /v1/chat/completions — optimize and forward to OpenAI.
18pub async fn openai_handler(
19    State(state): State<ProxyState>,
20    uri: Uri,
21    headers: HeaderMap,
22    body: String,
23) -> Response {
24    let config = state.config.shift_config("openai");
25    let base_url = &state.config.providers.openai;
26    let query = uri.query().map(|q| format!("?{}", q)).unwrap_or_default();
27    let target_url = format!("{}{}{}", base_url, uri.path(), query);
28
29    // Run SHIFT optimization pipeline on a blocking thread to avoid
30    // starving the async runtime during CPU-intensive image operations.
31    let start = std::time::Instant::now();
32    let body_clone = body.clone();
33    let optimization_result =
34        tokio::task::spawn_blocking(move || optimize_payload(&body_clone, &config)).await;
35
36    let (final_body, optimized) = match optimization_result {
37        Ok(Some((transformed_json, report))) => {
38            let duration_ms = start.elapsed().as_millis() as u64;
39
40            // Record session stats (in-memory)
41            state.session.record(&report);
42
43            // Record persistent stats with FULL token savings
44            let record = shift_preflight::stats::record_from_report(&report, "openai", duration_ms);
45            if let Err(e) = shift_preflight::stats::record_run(&record, None) {
46                tracing::warn!("failed to save stats: {}", e);
47            }
48
49            if state.config.verbose {
50                let saved = report.original_size.saturating_sub(report.transformed_size);
51                if saved > 0 {
52                    tracing::info!(
53                        "OpenAI: saved {:.1}KB ({} tokens)",
54                        saved as f64 / 1024.0,
55                        report.token_savings.openai_saved(),
56                    );
57                }
58            }
59
60            (transformed_json, true)
61        }
62        Ok(None) | Err(_) => (body, false),
63    };
64
65    if state.config.verbose && !optimized {
66        tracing::debug!("OpenAI: no optimization applied (passthrough)");
67    }
68
69    forward_request(
70        &state.http_client,
71        "POST",
72        &target_url,
73        &headers,
74        Some(final_body),
75    )
76    .await
77}