lean_ctx/proxy/
google.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.gemini_upstream.clone();
18    forward::forward_request(
19        State(state),
20        req,
21        &upstream,
22        "/",
23        compress_request_body,
24        "Gemini",
25        &["application/x-ndjson"],
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(contents) = doc.get_mut("contents").and_then(|c| c.as_array_mut()) {
42        for content in contents.iter_mut() {
43            if let Some(parts) = content.get_mut("parts").and_then(|p| p.as_array_mut()) {
44                for part in parts.iter_mut() {
45                    if let Some(func_resp) = part.get_mut("functionResponse") {
46                        if let Some(response) = func_resp.get_mut("response") {
47                            modified |= compress_string_field(response, "result");
48                            modified |= compress_string_field(response, "content");
49                        }
50                    }
51                }
52            }
53        }
54    }
55
56    if !modified {
57        return (body.to_vec(), original_size, original_size);
58    }
59
60    match serde_json::to_vec(&doc) {
61        Ok(compressed) => {
62            let compressed_size = compressed.len();
63            (compressed, original_size, compressed_size)
64        }
65        Err(_) => (body.to_vec(), original_size, original_size),
66    }
67}
68
69fn compress_string_field(obj: &mut Value, field: &str) -> bool {
70    if let Some(val) = obj
71        .get_mut(field)
72        .and_then(|v| v.as_str().map(String::from))
73    {
74        let compressed = compress_tool_result(&val, None);
75        if compressed.len() < val.len() {
76            obj[field] = Value::String(compressed);
77            return true;
78        }
79    }
80    false
81}
lean_ctx/proxy/google.rs

lean_ctx/proxy/
google.rs
