1use axum::{
2 body::Body,
3 extract::State,
4 http::{Request, StatusCode},
5 response::Response,
6};
7use serde_json::Value;
8
9use super::compress::compress_tool_result;
10use super::forward;
11use super::ProxyState;
12
13pub async fn handler(
14 State(state): State<ProxyState>,
15 req: Request<Body>,
16) -> Result<Response, StatusCode> {
17 let upstream = state.openai_upstream.clone();
18 forward::forward_request(
19 State(state),
20 req,
21 &upstream,
22 "/v1/chat/completions",
23 compress_request_body,
24 "OpenAI",
25 &[],
26 )
27 .await
28}
29
30fn compress_request_body(body: &[u8]) -> (Vec<u8>, usize, usize) {
31 let original_size = body.len();
32
33 let parsed: Value = match serde_json::from_slice(body) {
34 Ok(v) => v,
35 Err(_) => return (body.to_vec(), original_size, original_size),
36 };
37
38 let mut doc = parsed;
39 let mut modified = false;
40
41 if let Some(messages) = doc.get_mut("messages").and_then(|m| m.as_array_mut()) {
42 for msg in messages.iter_mut() {
43 let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
44 if role != "tool" {
45 continue;
46 }
47
48 if let Some(content) = msg
49 .get_mut("content")
50 .and_then(|c| c.as_str().map(String::from))
51 {
52 let compressed = compress_tool_result(&content, None);
53 if compressed.len() < content.len() {
54 msg["content"] = Value::String(compressed);
55 modified = true;
56 }
57 }
58 }
59 }
60
61 if !modified {
62 return (body.to_vec(), original_size, original_size);
63 }
64
65 match serde_json::to_vec(&doc) {
66 Ok(compressed) => {
67 let compressed_size = compressed.len();
68 (compressed, original_size, compressed_size)
69 }
70 Err(_) => (body.to_vec(), original_size, original_size),
71 }
72}