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
13const UPSTREAM: &str = "https://generativelanguage.googleapis.com";
14
15pub async fn handler(state: State<ProxyState>, req: Request<Body>) -> Result<Response, StatusCode> {
16    forward::forward_request(
17        state,
18        req,
19        UPSTREAM,
20        "/",
21        compress_request_body,
22        "Gemini",
23        &["application/x-ndjson"],
24    )
25    .await
26}
27
28fn compress_request_body(body: &[u8]) -> (Vec<u8>, usize, usize) {
29    let original_size = body.len();
30
31    let parsed: Value = match serde_json::from_slice(body) {
32        Ok(v) => v,
33        Err(_) => return (body.to_vec(), original_size, original_size),
34    };
35
36    let mut doc = parsed;
37    let mut modified = false;
38
39    if let Some(contents) = doc.get_mut("contents").and_then(|c| c.as_array_mut()) {
40        for content in contents.iter_mut() {
41            if let Some(parts) = content.get_mut("parts").and_then(|p| p.as_array_mut()) {
42                for part in parts.iter_mut() {
43                    if let Some(func_resp) = part.get_mut("functionResponse") {
44                        if let Some(response) = func_resp.get_mut("response") {
45                            modified |= compress_string_field(response, "result");
46                            modified |= compress_string_field(response, "content");
47                        }
48                    }
49                }
50            }
51        }
52    }
53
54    if !modified {
55        return (body.to_vec(), original_size, original_size);
56    }
57
58    match serde_json::to_vec(&doc) {
59        Ok(compressed) => {
60            let compressed_size = compressed.len();
61            (compressed, original_size, compressed_size)
62        }
63        Err(_) => (body.to_vec(), original_size, original_size),
64    }
65}
66
67fn compress_string_field(obj: &mut Value, field: &str) -> bool {
68    if let Some(val) = obj
69        .get_mut(field)
70        .and_then(|v| v.as_str().map(String::from))
71    {
72        let compressed = compress_tool_result(&val, None);
73        if compressed.len() < val.len() {
74            obj[field] = Value::String(compressed);
75            return true;
76        }
77    }
78    false
79}
lean_ctx/proxy/google.rs

lean_ctx/proxy/
google.rs
