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::tool_kind::{self, should_protect, ToolResultKind};
12use super::ProxyState;
13
14pub async fn handler(
15    State(state): State<ProxyState>,
16    req: Request<Body>,
17) -> Result<Response, StatusCode> {
18    let upstream = state.gemini_upstream.clone();
19    forward::forward_request(
20        State(state),
21        req,
22        &upstream,
23        "/",
24        compress_request_body,
25        "Gemini",
26        &["application/x-ndjson"],
27    )
28    .await
29}
30
31fn compress_request_body(parsed: Value, original_size: usize) -> (Vec<u8>, usize, usize) {
32    let mut doc = parsed;
33    let mut modified = false;
34
35    if let Some(contents) = doc.get_mut("contents").and_then(|c| c.as_array_mut()) {
36        for content in contents.iter_mut() {
37            if let Some(parts) = content.get_mut("parts").and_then(|p| p.as_array_mut()) {
38                for part in parts.iter_mut() {
39                    if let Some(func_resp) = part.get_mut("functionResponse") {
40                        // Gemini carries the originating function name inline.
41                        let kind = func_resp
42                            .get("name")
43                            .and_then(|v| v.as_str())
44                            .map_or(ToolResultKind::Other, tool_kind::classify_tool_name);
45                        if let Some(response) = func_resp.get_mut("response") {
46                            modified |= compress_string_field(response, "result", kind);
47                            modified |= compress_string_field(response, "content", kind);
48                        }
49                    }
50                }
51            }
52        }
53    }
54
55    let out = serde_json::to_vec(&doc).unwrap_or_default();
56    let compressed_size = if modified { out.len() } else { original_size };
57    (out, original_size, compressed_size)
58}
59
60fn compress_string_field(obj: &mut Value, field: &str, kind: ToolResultKind) -> bool {
61    if let Some(val) = obj
62        .get_mut(field)
63        .and_then(|v| v.as_str().map(String::from))
64    {
65        if should_protect(kind, &val) {
66            return false;
67        }
68        let compressed = compress_tool_result(&val, None);
69        if compressed.len() < val.len() {
70            obj[field] = Value::String(compressed);
71            return true;
72        }
73    }
74    false
75}
lean_ctx/proxy/google.rs

lean_ctx/proxy/
google.rs
