Skip to main content

lean_ctx/proxy/
anthropic.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.anthropic_upstream.clone();
18    forward::forward_request(
19        State(state),
20        req,
21        &upstream,
22        "/v1/messages",
23        compress_request_body,
24        "Anthropic",
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        for msg in messages.iter_mut() {
43            let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
44            if role != "user" {
45                continue;
46            }
47
48            if let Some(content) = msg.get_mut("content").and_then(|c| c.as_array_mut()) {
49                for block in content.iter_mut() {
50                    if block.get("type").and_then(|t| t.as_str()) != Some("tool_result") {
51                        continue;
52                    }
53
54                    if let Some(inner_content) = block.get_mut("content") {
55                        modified |= compress_content_field(inner_content, None);
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_content_field(content: &mut Value, tool_name: Option<&str>) -> bool {
76    match content {
77        Value::String(s) => {
78            let compressed = compress_tool_result(s, tool_name);
79            if compressed.len() < s.len() {
80                *s = compressed;
81                return true;
82            }
83            false
84        }
85        Value::Array(arr) => {
86            let mut modified = false;
87            for item in arr.iter_mut() {
88                if item.get("type").and_then(|t| t.as_str()) == Some("text") {
89                    if let Some(text) = item
90                        .get_mut("text")
91                        .and_then(|t| t.as_str().map(String::from))
92                    {
93                        let compressed = compress_tool_result(&text, tool_name);
94                        if compressed.len() < text.len() {
95                            item["text"] = Value::String(compressed);
96                            modified = true;
97                        }
98                    }
99                }
100            }
101            modified
102        }
103        _ => false,
104    }
105}