llmshim 0.11.0

Blazing fast LLM API translation layer in pure Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

#![cfg(feature = "proxy")]

use axum::{
    body::{to_bytes, Body},
    http::Request,
};
use llmshim::{
    client::ShimClient, error::ShimError, log::Logger, providers::gemini::Gemini, router::Router,
};
use mockito::Matcher;
use serde_json::json;
use std::error::Error as StdError;
use tower::ServiceExt;

const SYNTHETIC_GEMINI_KEY: &str = "test-gemini-key";
const GEMINI_MODEL: &str = "gemini-test-model";

fn request() -> serde_json::Value {
    json!({
        "model": format!("gemini/{GEMINI_MODEL}"),
        "messages": [{"role": "user", "content": "test"}],
    })
}

fn provider(base_url: String) -> Gemini {
    Gemini::new(SYNTHETIC_GEMINI_KEY.into()).with_base_url(base_url)
}

fn sensitive_url(base_url: &str) -> String {
    format!("{base_url}/models/{GEMINI_MODEL}:generateContent?key={SYNTHETIC_GEMINI_KEY}")
}

fn assert_error_has_no_sensitive_url(
    error: &ShimError,
    expected_sensitive_url: &str,
    error_kind: &str,
) {
    let rendered_error = error.to_string();
    let debug_error = format!("{error:?}");
    for rendered_value in [&rendered_error, &debug_error] {
        assert!(
            !rendered_value.contains(SYNTHETIC_GEMINI_KEY),
            "{error_kind} error retained the synthetic query credential"
        );
        assert!(
            !rendered_value.contains(expected_sensitive_url),
            "{error_kind} error retained the full request URL"
        );
    }
    let source = StdError::source(error).expect("HTTP errors retain their source");
    let rendered_source = source.to_string();
    assert!(
        !rendered_source.contains(SYNTHETIC_GEMINI_KEY),
        "{error_kind} error source retained the synthetic query credential"
    );
    assert!(
        !rendered_source.contains(expected_sensitive_url),
        "{error_kind} error source retained the full request URL"
    );
}

#[tokio::test]
async fn gemini_transport_and_malformed_json_errors_redact_query_credentials() {
    std::env::set_var("LLMSHIM_MAX_RETRIES", "0");

    let closed_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let closed_base_url = format!("http://{}", closed_listener.local_addr().unwrap());
    drop(closed_listener);

    let transport_error = ShimClient::new()
        .completion(&provider(closed_base_url.clone()), GEMINI_MODEL, &request())
        .await
        .expect_err("a closed loopback port must produce a transport error");
    assert!(matches!(transport_error, ShimError::Http(_)));
    assert_error_has_no_sensitive_url(
        &transport_error,
        &sensitive_url(&closed_base_url),
        "transport",
    );

    let mut upstream_server = mockito::Server::new_async().await;
    let malformed_response = upstream_server
        .mock("POST", "/models/gemini-test-model:generateContent")
        .match_query(Matcher::UrlEncoded(
            "key".into(),
            SYNTHETIC_GEMINI_KEY.into(),
        ))
        .with_status(200)
        .with_body("this is not JSON")
        .expect(2)
        .create_async()
        .await;
    let malformed_base_url = upstream_server.url();
    let expected_sensitive_url = sensitive_url(&malformed_base_url);

    let malformed_json_error = ShimClient::new()
        .completion(
            &provider(malformed_base_url.clone()),
            GEMINI_MODEL,
            &request(),
        )
        .await
        .expect_err("a malformed successful response must fail JSON decoding");
    assert!(matches!(malformed_json_error, ShimError::Http(_)));
    assert_error_has_no_sensitive_url(
        &malformed_json_error,
        &expected_sensitive_url,
        "malformed JSON",
    );

    let log_directory = tempfile::tempdir().unwrap();
    let log_path = log_directory.path().join("proxy.jsonl");
    let logger = Logger::to_file(log_path.to_str().unwrap()).unwrap();
    let router = Router::new().register("gemini", Box::new(provider(malformed_base_url.clone())));
    let response = llmshim::proxy::app(router, Some(logger))
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/v1/chat")
                .header("content-type", "application/json")
                .body(Body::from(request().to_string()))
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(response.status(), 502);
    let response_body = to_bytes(response.into_body(), 100_000).await.unwrap();
    let rendered_response = String::from_utf8(response_body.to_vec()).unwrap();
    assert!(!rendered_response.contains(SYNTHETIC_GEMINI_KEY));
    assert!(!rendered_response.contains(&expected_sensitive_url));

    let rendered_log = std::fs::read_to_string(&log_path).unwrap();
    assert!(!rendered_log.contains(SYNTHETIC_GEMINI_KEY));
    assert!(!rendered_log.contains(&expected_sensitive_url));
    malformed_response.assert_async().await;
}