File size: 2,107 Bytes
854994d | 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 | use crate::common::SafetyBufferingTreatment;
use http::HeaderMap;
pub(crate) const X_CODEX_SAFETY_BUFFERING_ENABLED_HEADER: &str = "x-codex-safety-buffering-enabled";
pub(crate) const X_CODEX_SAFETY_BUFFERING_FASTER_MODEL_HEADER: &str =
"x-codex-safety-buffering-faster-model";
pub(crate) fn treatment_from_headers(headers: &HeaderMap) -> Option<SafetyBufferingTreatment> {
if !headers.contains_key(X_CODEX_SAFETY_BUFFERING_ENABLED_HEADER)
&& !headers.contains_key(X_CODEX_SAFETY_BUFFERING_FASTER_MODEL_HEADER)
{
return None;
}
let faster_model = headers
.get(X_CODEX_SAFETY_BUFFERING_FASTER_MODEL_HEADER)
.and_then(|value| value.to_str().ok())
.map(str::to_string);
Some(SafetyBufferingTreatment { faster_model })
}
#[cfg(test)]
mod tests {
use super::*;
use http::HeaderValue;
use pretty_assertions::assert_eq;
#[test]
fn reads_treatment_from_http_headers() {
let mut headers = HeaderMap::new();
headers.insert(
X_CODEX_SAFETY_BUFFERING_ENABLED_HEADER,
HeaderValue::from_static("true"),
);
headers.insert(
X_CODEX_SAFETY_BUFFERING_FASTER_MODEL_HEADER,
HeaderValue::from_static("faster-model"),
);
assert_eq!(
treatment_from_headers(&headers),
Some(SafetyBufferingTreatment {
faster_model: Some("faster-model".to_string()),
})
);
}
#[test]
fn buffering_enabled_header_does_not_gate_the_faster_model_fallback() {
let mut headers = HeaderMap::new();
headers.insert(
X_CODEX_SAFETY_BUFFERING_ENABLED_HEADER,
HeaderValue::from_static("false"),
);
headers.insert(
X_CODEX_SAFETY_BUFFERING_FASTER_MODEL_HEADER,
HeaderValue::from_static("faster-model"),
);
assert_eq!(
treatment_from_headers(&headers),
Some(SafetyBufferingTreatment {
faster_model: Some("faster-model".to_string()),
})
);
}
}
|