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
13const UPSTREAM: &str = "https://api.anthropic.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        "/v1/messages",
21        compress_request_body,
22        "Anthropic",
23        &[],
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(messages) = doc.get_mut("messages").and_then(|m| m.as_array_mut()) {
40        for msg in messages.iter_mut() {
41            let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
42            if role != "user" {
43                continue;
44            }
45
46            if let Some(content) = msg.get_mut("content").and_then(|c| c.as_array_mut()) {
47                for block in content.iter_mut() {
48                    if block.get("type").and_then(|t| t.as_str()) != Some("tool_result") {
49                        continue;
50                    }
51
52                    if let Some(inner_content) = block.get_mut("content") {
53                        modified |= compress_content_field(inner_content, None);
54                    }
55                }
56            }
57        }
58    }
59
60    if !modified {
61        return (body.to_vec(), original_size, original_size);
62    }
63
64    match serde_json::to_vec(&doc) {
65        Ok(compressed) => {
66            let compressed_size = compressed.len();
67            (compressed, original_size, compressed_size)
68        }
69        Err(_) => (body.to_vec(), original_size, original_size),
70    }
71}
72
73fn compress_content_field(content: &mut Value, tool_name: Option<&str>) -> bool {
74    match content {
75        Value::String(s) => {
76            let compressed = compress_tool_result(s, tool_name);
77            if compressed.len() < s.len() {
78                *s = compressed;
79                return true;
80            }
81            false
82        }
83        Value::Array(arr) => {
84            let mut modified = false;
85            for item in arr.iter_mut() {
86                if item.get("type").and_then(|t| t.as_str()) == Some("text") {
87                    if let Some(text) = item
88                        .get_mut("text")
89                        .and_then(|t| t.as_str().map(String::from))
90                    {
91                        let compressed = compress_tool_result(&text, tool_name);
92                        if compressed.len() < text.len() {
93                            item["text"] = Value::String(compressed);
94                            modified = true;
95                        }
96                    }
97                }
98            }
99            modified
100        }
101        _ => false,
102    }
103}