wiremux 0.10.2

Wiremux maps Chat Completions, Messages, Responses, Gemini, and Converse through one IR.
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

//! Join profile `base_url` + `chat_path` (Gemini `{model}` + stream rewrite).

use wiremux_auth::{ResolvedProfile, Wire};

/// Shared miss text for `list_models` / `upstream_url_for_model`.
pub(crate) const MISSING_BASE_URL: &str =
    "profile has no base_url; set top-level `base_url` on the profile";

/// Join `base_url` + `chat_path`, substituting `{model}` when present.
///
/// Gemini streaming uses `:streamGenerateContent?alt=sse` when the path
/// is the unary `:generateContent` default.
pub fn upstream_url_for_model(
    profile: &ResolvedProfile,
    model: Option<&str>,
    stream: bool,
) -> Result<String, String> {
    let base = profile.http.base_url.as_deref().ok_or(MISSING_BASE_URL)?;
    let path = profile
        .http
        .chat_path
        .as_deref()
        .or_else(|| profile.dialect.wire.map(Wire::default_chat_path))
        .unwrap_or("/");
    let mut path = if let Some(model) = model.filter(|m| !m.is_empty()) {
        path.replace("{model}", &encode_model_segment(model))
    } else {
        path.to_string()
    };
    if path.contains("{model}") {
        return Err("request has no model and the profile path still contains {model}".to_string());
    }
    if stream
        && matches!(profile.dialect.wire, Some(Wire::Gemini))
        && path.ends_with(":generateContent")
    {
        path = path.replacen(":generateContent", ":streamGenerateContent?alt=sse", 1);
    }
    if stream
        && matches!(profile.dialect.wire, Some(Wire::Messages))
        && path.ends_with(":rawPredict")
    {
        path = path.replacen(":rawPredict", ":streamRawPredict", 1);
    }
    if stream && matches!(profile.dialect.wire, Some(Wire::Converse)) && path.ends_with("/converse")
    {
        path = format!("{path}-stream");
    }
    if path.starts_with("http://") || path.starts_with("https://") {
        return Ok(rewrite_vertex_global_host(path));
    }
    Ok(rewrite_vertex_global_host(format!(
        "{}{}",
        base.trim_end_matches('/'),
        if path.starts_with('/') {
            path
        } else {
            format!("/{path}")
        }
    )))
}

/// Encode `{model}` as one path segment. Keep `:` so Bedrock model ids
/// stay readable. Encode `/ ? # % \\` and controls so the name cannot
/// change the path, query, or fragment.
pub(crate) fn encode_model_segment(model: &str) -> String {
    let mut out = String::with_capacity(model.len());
    for byte in model.bytes() {
        match byte {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' | b':' => {
                out.push(byte as char)
            }
            _ => {
                out.push('%');
                out.push(char::from(b"0123456789ABCDEF"[(byte >> 4) as usize]));
                out.push(char::from(b"0123456789ABCDEF"[(byte & 0x0F) as usize]));
            }
        }
    }
    out
}

/// Ingest writes `https://{env:GOOGLE_VERTEX_LOCATION}-aiplatform.googleapis.com`.
/// After envsubst, `LOCATION=global` becomes `global-aiplatform.googleapis.com`.
/// Google's global host is `aiplatform.googleapis.com` (path keeps `/locations/global/`).
fn rewrite_vertex_global_host(url: String) -> String {
    url.replacen(
        "://global-aiplatform.googleapis.com",
        "://aiplatform.googleapis.com",
        1,
    )
}