Skip to main content

lean_ctx/proxy/
openai.rs

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        super::history_prune::prune_history(messages, 6);
43        modified = true;
44
45        for msg in messages.iter_mut() {
46            let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
47            if role != "tool" {
48                continue;
49            }
50
51            if let Some(content) = msg
52                .get_mut("content")
53                .and_then(|c| c.as_str().map(String::from))
54            {
55                let compressed = compress_tool_result(&content, None);
56                if compressed.len() < content.len() {
57                    msg["content"] = Value::String(compressed);
58                    modified = true;
59                }
60            }
61        }
62    }
63
64    if !modified {
65        return (body.to_vec(), original_size, original_size);
66    }
67
68    match serde_json::to_vec(&doc) {
69        Ok(compressed) => {
70            let compressed_size = compressed.len();
71            (compressed, original_size, compressed_size)
72        }
73        Err(_) => (body.to_vec(), original_size, original_size),
74    }
75}