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(body: &[u8]) -> (Vec<u8>, usize, usize) {
32    let original_size = body.len();
33
34    let parsed: Value = match serde_json::from_slice(body) {
35        Ok(v) => v,
36        Err(_) => return (body.to_vec(), original_size, original_size),
37    };
38
39    let mut doc = parsed;
40    let mut modified = false;
41
42    if let Some(contents) = doc.get_mut("contents").and_then(|c| c.as_array_mut()) {
43        for content in contents.iter_mut() {
44            if let Some(parts) = content.get_mut("parts").and_then(|p| p.as_array_mut()) {
45                for part in parts.iter_mut() {
46                    if let Some(func_resp) = part.get_mut("functionResponse") {
47                        // Gemini carries the originating function name inline.
48                        let kind = func_resp
49                            .get("name")
50                            .and_then(|v| v.as_str())
51                            .map_or(ToolResultKind::Other, tool_kind::classify_tool_name);
52                        if let Some(response) = func_resp.get_mut("response") {
53                            modified |= compress_string_field(response, "result", kind);
54                            modified |= compress_string_field(response, "content", kind);
55                        }
56                    }
57                }
58            }
59        }
60    }
61
62    if !modified {
63        return (body.to_vec(), original_size, original_size);
64    }
65
66    match serde_json::to_vec(&doc) {
67        Ok(compressed) => {
68            let compressed_size = compressed.len();
69            (compressed, original_size, compressed_size)
70        }
71        Err(_) => (body.to_vec(), original_size, original_size),
72    }
73}
74
75fn compress_string_field(obj: &mut Value, field: &str, kind: ToolResultKind) -> bool {
76    if let Some(val) = obj
77        .get_mut(field)
78        .and_then(|v| v.as_str().map(String::from))
79    {
80        if should_protect(kind, &val) {
81            return false;
82        }
83        let compressed = compress_tool_result(&val, None);
84        if compressed.len() < val.len() {
85            obj[field] = Value::String(compressed);
86            return true;
87        }
88    }
89    false
90}
lean_ctx/proxy/google.rs

lean_ctx/proxy/
google.rs
