mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge c548955e9d into e84859ae4e
This commit is contained in:
Nate Armstrong
GitHub
commit
bcf9cdc7fd
44 changed files with 2446 additions and 86 deletions
1
litellm-rust/Cargo.lock
generated
1
litellm-rust/Cargo.lock
generated
|
|
@ -4320,6 +4320,7 @@ dependencies = [
|
|||
"thiserror 2.0.19",
|
||||
"time",
|
||||
"tokio",
|
||||
"unicode-general-category",
|
||||
"url",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -4,8 +4,10 @@ use axum::{
|
|||
body::{Body, to_bytes},
|
||||
http::Request,
|
||||
};
|
||||
use futures_util::{StreamExt, TryStreamExt};
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tower::ServiceExt;
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
|
|
@ -75,25 +77,76 @@ async fn messages_reaches_the_provider_and_preserves_json_or_sse(
|
|||
assert!(!requests[0].headers.contains_key("x-request-id"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::invalid_json("{".into())]
|
||||
#[case::array("[]".into())]
|
||||
#[case::null("null".into())]
|
||||
#[case::missing_model(json!({"messages": [], "max_tokens": 16}).to_string())]
|
||||
#[case::unknown_model(json!({"model": "missing/model", "messages": [], "max_tokens": 16}).to_string())]
|
||||
#[case::invalid_messages(json!({"model": "public/model", "messages": "invalid", "max_tokens": 16}).to_string())]
|
||||
#[case::missing_max_tokens(json!({"model": "public/model", "messages": [{"role": "user", "content": "hello"}]}).to_string())]
|
||||
#[tokio::test]
|
||||
async fn invalid_messages_stays_an_anthropic_error() {
|
||||
let response = support::post(
|
||||
support::app("anthropic/test-model", "http://127.0.0.1:1"),
|
||||
"/v1/messages",
|
||||
json!({"model": "public/model", "messages": "invalid", "max_tokens": 16}),
|
||||
)
|
||||
.await;
|
||||
async fn invalid_messages_stays_an_anthropic_error_without_provider_io(#[case] body: String) {
|
||||
let upstream = MockServer::start().await;
|
||||
let request = Request::post("/v1/messages")
|
||||
.header("content-type", "application/json")
|
||||
.header("x-request-id", "invalid-messages-request")
|
||||
.body(Body::from(body))
|
||||
.unwrap();
|
||||
let response = support::app("anthropic/test-model", &upstream.uri())
|
||||
.oneshot(request)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), 400);
|
||||
let body = support::json(response).await;
|
||||
assert_eq!(body["type"], "error");
|
||||
assert_eq!(body["error"]["type"], "invalid_request_error");
|
||||
assert_eq!(body["request_id"], "invalid-messages-request");
|
||||
assert!(!body["error"]["message"].as_str().unwrap().is_empty());
|
||||
assert!(upstream.received_requests().await.unwrap().is_empty());
|
||||
}
|
||||
|
||||
const FIRST_EVENT: &str = "event: message_start\ndata: {}\n\n";
|
||||
|
||||
async fn interrupted_stream(
|
||||
prefix: &'static [u8],
|
||||
) -> (
|
||||
String,
|
||||
tokio::sync::oneshot::Sender<()>,
|
||||
tokio::task::JoinHandle<()>,
|
||||
) {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base = format!("http://{}", listener.local_addr().unwrap());
|
||||
let (release, released) = tokio::sync::oneshot::channel();
|
||||
let task = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
let mut request = vec![0; 4096];
|
||||
assert!(socket.read(&mut request).await.unwrap() > 0);
|
||||
socket.write_all(
|
||||
format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n",
|
||||
prefix.len() + 1024
|
||||
).as_bytes()
|
||||
).await.unwrap();
|
||||
socket.write_all(prefix).await.unwrap();
|
||||
let _ = released.await;
|
||||
});
|
||||
(base, release, task)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::complete_lf(FIRST_EVENT.as_bytes())]
|
||||
#[case::complete_crlf(b"event: message_start\r\ndata: {}\r\n\r\n")]
|
||||
#[case::partial_event(b"event: content_blo")]
|
||||
#[case::partial_data(b"event: content_block_delta\ndata: {\"type\":\"content_block_delta\"")]
|
||||
#[case::partial_string(b"event: content_block_delta\ndata: {\"delta\":{\"text\":\"partial")]
|
||||
#[case::partial_utf8(b"event: content_block_delta\ndata: {\"delta\":{\"text\":\"\xe9")]
|
||||
#[tokio::test]
|
||||
async fn a_stream_that_fails_after_opening_ends_with_an_sse_error_frame() {
|
||||
let base = support::truncating_upstream(FIRST_EVENT).await;
|
||||
async fn a_stream_that_fails_after_opening_ends_with_an_sse_error_frame(
|
||||
#[case] prefix: &'static [u8],
|
||||
) {
|
||||
let (base, release, upstream) = interrupted_stream(prefix).await;
|
||||
let request = Request::post("/v1/messages")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
|
|
@ -109,12 +162,26 @@ async fn a_stream_that_fails_after_opening_ends_with_an_sse_error_frame() {
|
|||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), 200);
|
||||
let body = to_bytes(response.into_body(), 4096).await.unwrap();
|
||||
let text = std::str::from_utf8(&body).unwrap();
|
||||
let mut chunks = response.into_body().into_data_stream();
|
||||
let mut delivered = Vec::with_capacity(prefix.len());
|
||||
while delivered.len() < prefix.len() {
|
||||
let chunk = tokio::time::timeout(std::time::Duration::from_secs(5), chunks.next())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
delivered.extend_from_slice(&chunk);
|
||||
}
|
||||
assert_eq!(delivered, prefix);
|
||||
release.send(()).unwrap();
|
||||
let tail = chunks.try_collect::<Vec<_>>().await.unwrap().concat();
|
||||
upstream.await.unwrap();
|
||||
let text = std::str::from_utf8(&tail).unwrap();
|
||||
let frame = text
|
||||
.strip_prefix(FIRST_EVENT)
|
||||
.and_then(|rest| rest.strip_prefix("event: error\ndata: "))
|
||||
.unwrap_or_else(|| panic!("the delivered event then one error frame, got {text:?}"));
|
||||
.strip_prefix("\n\nevent: error\ndata: ")
|
||||
.unwrap_or_else(|| {
|
||||
panic!("an independent error frame after the delivered prefix, got {text:?}")
|
||||
});
|
||||
let error: serde_json::Value = serde_json::from_str(frame.trim_end()).unwrap();
|
||||
assert_eq!(error["type"], "error");
|
||||
assert_eq!(error["error"]["type"], "api_error");
|
||||
|
|
@ -217,3 +284,53 @@ async fn openrouter_receives_native_messages_headers_and_its_deployment_bearer(
|
|||
assert_eq!(requests.len(), 1);
|
||||
assert!(!requests[0].headers.contains_key("x-api-key"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::u64_max("18446744073709551615".to_string(), 200)]
|
||||
#[case::above_u64("18446744073709551617".to_string(), 500)]
|
||||
#[case::i64_min("-9223372036854775808".to_string(), 200)]
|
||||
#[case::below_i64("-9223372036854775809".to_string(), 500)]
|
||||
#[case::long_numeric_string(format!(r#""{}""#, "9".repeat(4096)), 200)]
|
||||
#[case::decimal("18446744073709551617.0".to_string(), 200)]
|
||||
#[case::exponent("1e30".to_string(), 200)]
|
||||
#[tokio::test]
|
||||
async fn messages_response_integer_overflow_is_a_provider_error_not_rounded_success(
|
||||
#[case] token: String,
|
||||
#[case] status: u16,
|
||||
) {
|
||||
let upstream = MockServer::start().await;
|
||||
let body = format!(
|
||||
r#"{{"id":"m","type":"message","role":"assistant","model":"test-model","content":[],"future_extension":{{"wide_integer":{token}}}}}"#,
|
||||
);
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(body, "application/json"))
|
||||
.expect(1)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let request = Request::post("/v1/messages")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(json!({"model":"public/model","messages":[{"role":"user","content":"hi"}],"max_tokens":16}).to_string()))
|
||||
.unwrap();
|
||||
let response = support::app("anthropic/test-model", &upstream.uri())
|
||||
.oneshot(request)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), status);
|
||||
let response = support::json(response).await;
|
||||
if status == 200 {
|
||||
assert_eq!(
|
||||
response["future_extension"]["wide_integer"],
|
||||
serde_json::from_str::<serde_json::Value>(&token).unwrap()
|
||||
);
|
||||
} else {
|
||||
assert_eq!(response["type"], "error");
|
||||
assert_eq!(response["error"]["type"], "api_error");
|
||||
assert!(
|
||||
response["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("response integer outside")
|
||||
);
|
||||
}
|
||||
assert_eq!(upstream.received_requests().await.unwrap().len(), 1);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -134,7 +134,7 @@ async fn a_stream_that_fails_after_opening_ends_with_the_responses_error_event()
|
|||
let text = std::str::from_utf8(&body).unwrap();
|
||||
let frame = text
|
||||
.strip_prefix(RESPONSE_CREATED)
|
||||
.and_then(|rest| rest.strip_prefix("event: error\ndata: "))
|
||||
.and_then(|rest| rest.strip_prefix("\n\nevent: error\ndata: "))
|
||||
.unwrap_or_else(|| panic!("the delivered event then one error frame, got {text:?}"));
|
||||
let error: serde_json::Value = serde_json::from_str(frame.trim_end()).unwrap();
|
||||
assert_eq!(error["type"], "error");
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ where
|
|||
}
|
||||
|
||||
fn encode_stream_error(&self, error: Error<P::Error>) -> Bytes {
|
||||
(self.stream_error)(error)
|
||||
let frame = (self.stream_error)(error);
|
||||
Bytes::from([b"\n\n".as_slice(), frame.as_ref()].concat())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -40,11 +40,15 @@ impl Protocol for TestProtocol {
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::complete(false)]
|
||||
#[case::failed(true)]
|
||||
#[case::complete(false, "event: custom\ndata: first\n\n")]
|
||||
#[case::failed_after_frame(true, "event: custom\ndata: first\n\n")]
|
||||
#[case::failed_during_frame(true, "event: custom\ndata: partial")]
|
||||
#[tokio::test]
|
||||
async fn sse_preserves_encoded_chunks_and_uses_the_supplied_error_format(#[case] fail: bool) {
|
||||
let first = Bytes::from_static(b"event: custom\ndata: first\n\n");
|
||||
async fn sse_preserves_encoded_chunks_and_delimits_the_supplied_error_frame(
|
||||
#[case] fail: bool,
|
||||
#[case] prefix: &'static str,
|
||||
) {
|
||||
let first = Bytes::from_static(prefix.as_bytes());
|
||||
let last = Bytes::from_static(b"data: [DONE]\n\n");
|
||||
let machine =
|
||||
hosted_call::<TestProtocol, _, _>((), None, move |(), _, _, _observations| async move {
|
||||
|
|
@ -70,14 +74,12 @@ async fn sse_preserves_encoded_chunks_and_uses_the_supplied_error_format(#[case]
|
|||
assert_eq!(response.headers()[CONTENT_TYPE], "text/event-stream");
|
||||
assert_eq!(errors.load(Ordering::SeqCst), 0);
|
||||
let body = to_bytes(response.into_body(), 1024).await.unwrap();
|
||||
assert_eq!(
|
||||
body,
|
||||
if fail {
|
||||
"event: custom\ndata: first\n\nevent: custom_error\ndata: Call(Upstream)\n\n"
|
||||
} else {
|
||||
"event: custom\ndata: first\n\ndata: [DONE]\n\n"
|
||||
}
|
||||
);
|
||||
let expected = if fail {
|
||||
format!("{prefix}\n\nevent: custom_error\ndata: Call(Upstream)\n\n")
|
||||
} else {
|
||||
format!("{prefix}data: [DONE]\n\n")
|
||||
};
|
||||
assert_eq!(body.as_ref(), expected.as_bytes());
|
||||
assert_eq!(errors.load(Ordering::SeqCst), usize::from(fail));
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ litellm-secrets.workspace = true
|
|||
litellm-tracing.workspace = true
|
||||
reqwest.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json = { workspace = true, features = ["preserve_order"] }
|
||||
serde_json = { workspace = true, features = ["preserve_order", "raw_value"] }
|
||||
tokio.workspace = true
|
||||
tracing.workspace = true
|
||||
|
||||
|
|
|
|||
|
|
@ -72,6 +72,12 @@ pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option<M
|
|||
Some(MessagesProvider { provider, config })
|
||||
}
|
||||
|
||||
pub fn supports_messages_provider(model: &str, custom_llm_provider: Option<&str>) -> bool {
|
||||
litellm_inference::provider::resolve_llm_provider(model, custom_llm_provider, "messages")
|
||||
.ok()
|
||||
.is_some_and(|provider| messages_provider(provider.provider, provider.model).is_some())
|
||||
}
|
||||
|
||||
pub(super) fn string_headers(
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
|
|
@ -84,7 +90,7 @@ mod tests {
|
|||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{messages_provider, string_headers};
|
||||
use super::{messages_provider, string_headers, supports_messages_provider};
|
||||
use crate::Error;
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
|
||||
|
|
@ -113,6 +119,19 @@ mod tests {
|
|||
assert!(messages_provider(provider, model).is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::anthropic("anthropic/claude-test", None, true)]
|
||||
#[case::inferred_bare_model("claude-test", Some("anthropic"), true)]
|
||||
#[case::generic_passthrough("openai/gpt-5", None, false)]
|
||||
#[case::vertex_non_claude("vertex_ai/gemini", None, false)]
|
||||
fn admission_requires_a_native_adapter(
|
||||
#[case] model: &str,
|
||||
#[case] provider: Option<&str>,
|
||||
#[case] supported: bool,
|
||||
) {
|
||||
assert_eq!(supports_messages_provider(model, provider), supported);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn string_headers_rejects_non_string_values() {
|
||||
let headers = json!({"x-count": 3}).as_object().unwrap().clone();
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use std::time::Duration;
|
||||
use std::{collections::BTreeMap, time::Duration};
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures_util::{FutureExt, StreamExt, TryStreamExt, stream::BoxStream};
|
||||
|
|
@ -77,7 +77,10 @@ pub(super) async fn execute(
|
|||
&request.wire,
|
||||
);
|
||||
let identity = request.identity.clone();
|
||||
let (output, source) = match cache.lookup().await {
|
||||
let hit = cache.lookup().await.filter(|(output, _)| {
|
||||
matches!(output, MessagesCallResponse::Stream { .. }) == request.stream
|
||||
});
|
||||
let (output, source) = match hit {
|
||||
Some(hit) => hit,
|
||||
None => (
|
||||
call_provider(&route.http, request, context).await?,
|
||||
|
|
@ -412,18 +415,63 @@ fn decode_response(
|
|||
model: &str,
|
||||
text: &str,
|
||||
) -> Result<MessagesResponse, Error> {
|
||||
let response = serde_json::from_str(text).map_err(|err| {
|
||||
Error::invalid_response(litellm_llms::ErrorDetail::invalid(
|
||||
"messages response JSON",
|
||||
err,
|
||||
))
|
||||
})?;
|
||||
validate_response_numbers(text)?;
|
||||
let response = serde_json::from_str(text).map_err(invalid_response_json)?;
|
||||
config
|
||||
.transform_anthropic_messages_response(model, response)
|
||||
.map_err(Error::from)
|
||||
.map_err(Error::after_send)
|
||||
}
|
||||
|
||||
fn invalid_response_json(error: serde_json::Error) -> Error {
|
||||
Error::invalid_response(litellm_llms::ErrorDetail::invalid(
|
||||
"messages response JSON",
|
||||
error,
|
||||
))
|
||||
}
|
||||
|
||||
fn validate_response_numbers(text: &str) -> Result<(), Error> {
|
||||
use serde_json::value::RawValue;
|
||||
|
||||
let root: &RawValue = serde_json::from_str(text).map_err(invalid_response_json)?;
|
||||
let mut pending = vec![(root, 0usize)];
|
||||
while let Some((value, depth)) = pending.pop() {
|
||||
let token = value.get().trim();
|
||||
if depth >= 128 {
|
||||
return Err(invalid_response_json(serde::de::Error::custom(
|
||||
"messages response nesting exceeds 128 levels",
|
||||
)));
|
||||
}
|
||||
match token.as_bytes().first() {
|
||||
Some(b'{') => {
|
||||
let fields: BTreeMap<String, &RawValue> =
|
||||
serde_json::from_str(token).map_err(invalid_response_json)?;
|
||||
pending.extend(fields.into_values().map(|value| (value, depth + 1)));
|
||||
}
|
||||
Some(b'[') => {
|
||||
let items: Vec<&RawValue> =
|
||||
serde_json::from_str(token).map_err(invalid_response_json)?;
|
||||
pending.extend(items.into_iter().map(|value| (value, depth + 1)));
|
||||
}
|
||||
Some(b'-' | b'0'..=b'9') if !token.contains(['.', 'e', 'E']) => {
|
||||
let integer = if token.starts_with('-') {
|
||||
token.parse::<i64>().map(|_| ())
|
||||
} else {
|
||||
token.parse::<u64>().map(|_| ())
|
||||
};
|
||||
integer.map_err(|error| {
|
||||
Error::invalid_response(litellm_llms::ErrorDetail::invalid(
|
||||
"messages response integer outside i64/u64 range",
|
||||
error,
|
||||
))
|
||||
})?;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn streaming_response(
|
||||
response: reqwest::Response,
|
||||
decoder: Option<StreamDecoder>,
|
||||
|
|
@ -495,6 +543,46 @@ mod tests {
|
|||
use crate::{MessagesSettings, MessagesShaping, common_utils::messages_provider};
|
||||
use litellm_inference::provider::resolve_llm_provider;
|
||||
|
||||
#[rstest]
|
||||
#[case::u64_max("18446744073709551615", Some(json!(u64::MAX)))]
|
||||
#[case::above_u64("18446744073709551617", None)]
|
||||
#[case::i64_min("-9223372036854775808", Some(json!(i64::MIN)))]
|
||||
#[case::below_i64("-9223372036854775809", None)]
|
||||
#[case::decimal("18446744073709551617.0", Some(json!(18446744073709551617.0_f64)))]
|
||||
#[case::exponent("1.8446744073709552e19", Some(json!(1.8446744073709552e19_f64)))]
|
||||
#[case::quoted_integer(r#""18446744073709551617""#, Some(json!("18446744073709551617")))]
|
||||
#[case::escaped_digit(r#""\u00318446744073709551617""#, Some(json!("18446744073709551617")))]
|
||||
fn response_integer_tokens_never_round_into_success(
|
||||
#[case] token: &str,
|
||||
#[case] expected: Option<Value>,
|
||||
) {
|
||||
let text = format!(
|
||||
r#"{{"id":"m","type":"message","role":"assistant","model":"claude-test","content":[],"future":{{"values":[{token}]}}}}"#,
|
||||
);
|
||||
let config = litellm_llms::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
|
||||
let result = decode_response(&config, "claude-test", &text);
|
||||
match expected {
|
||||
Some(expected) => assert_eq!(result.unwrap().extra["future"]["values"][0], expected),
|
||||
None => {
|
||||
let error = result.unwrap_err();
|
||||
assert!(error.is_response());
|
||||
assert!(!error.is_request());
|
||||
assert!(error.dispatched());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::long_numeric_string(format!(r#""{}""#, "9".repeat(4096)), true)]
|
||||
#[case::long_integer("9".repeat(4096), false)]
|
||||
#[case::deeply_nested(format!("{}0{}", "[".repeat(512), "]".repeat(512)), false)]
|
||||
fn response_numeric_guard_bounds_nesting_and_ignores_string_contents(
|
||||
#[case] token: String,
|
||||
#[case] accepted: bool,
|
||||
) {
|
||||
assert_eq!(validate_response_numbers(&token).is_ok(), accepted);
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn shaping() -> MessagesShaping {
|
||||
MessagesShaping::default()
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ use std::sync::Arc;
|
|||
|
||||
use crate::common_utils::messages_provider;
|
||||
|
||||
pub use common_utils::merge_messages_header_sources;
|
||||
pub use common_utils::{merge_messages_header_sources, supports_messages_provider};
|
||||
pub use litellm_inference::Error;
|
||||
pub use types::{
|
||||
MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params,
|
||||
|
|
|
|||
|
|
@ -293,3 +293,285 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_
|
|||
assert!(matches!(cached.source, ResultSource::Cache { .. }));
|
||||
assert_eq!(cached.cost, Cost::CacheHit);
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn streaming_call() -> litellm_inference_messages::MessagesCall {
|
||||
litellm_inference_messages::MessagesCall {
|
||||
body: serde_json::from_value(json!({
|
||||
"model": "anthropic/cache-test-model",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 32,
|
||||
"stream": true
|
||||
}))
|
||||
.unwrap(),
|
||||
api_key: Some("stream-key".into()),
|
||||
api_base: None,
|
||||
custom_llm_provider: None,
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
shaping: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
const COMPLETE_STREAM: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
|
||||
|
||||
async fn streamed_bytes(output: litellm_inference_messages::MessagesCallResponse) -> Vec<u8> {
|
||||
use futures_util::TryStreamExt;
|
||||
let litellm_inference_messages::MessagesCallResponse::Stream { chunks, .. } = output else {
|
||||
panic!("expected a Messages stream");
|
||||
};
|
||||
chunks.try_collect::<Vec<_>>().await.unwrap().concat()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::complete(COMPLETE_STREAM.to_owned(), true)]
|
||||
#[case::data_only("data: {\"type\":\"message_stop\"}\n\n".into(), true)]
|
||||
#[case::crlf("event: message_stop\r\ndata: {\"type\":\"message_stop\"}\r\n\r\n".into(), true)]
|
||||
#[case::empty(String::new(), false)]
|
||||
#[case::missing_stop("data: {\"type\":\"message_start\"}\n\n".into(), false)]
|
||||
#[case::unterminated_stop("data: {\"type\":\"message_stop\"}".into(), false)]
|
||||
#[case::partial_tail(format!("{COMPLETE_STREAM}data: {{"), false)]
|
||||
#[case::error_before_stop(format!("data: {{\"type\":\"error\",\"error\":{{\"message\":\"failed\"}}}}\n\n{COMPLETE_STREAM}"), false)]
|
||||
#[case::error_after_stop(format!("{COMPLETE_STREAM}data: {{\"type\":\"error\"}}\n\n"), false)]
|
||||
#[case::invalid_json(format!("data: invalid\n\n{COMPLETE_STREAM}"), false)]
|
||||
#[case::event_name_without_type("event: message_stop\ndata: {}\n\n".into(), false)]
|
||||
#[case::over_entry_limit(format!("data: {{\"type\":\"content_block_delta\",\"text\":\"{}\"}}\n\n{COMPLETE_STREAM}", "x".repeat(4096)), false)]
|
||||
#[tokio::test]
|
||||
async fn stream_cache_eligibility_preserves_every_delivered_byte(
|
||||
cache: Arc<dyn ResponseCacheService>,
|
||||
streaming_call: litellm_inference_messages::MessagesCall,
|
||||
#[case] wire: String,
|
||||
#[case] cache_hit: bool,
|
||||
) {
|
||||
use wiremock::ResponseTemplate;
|
||||
let response = ResponseTemplate::new(200).set_body_raw(wire.clone(), "text/event-stream");
|
||||
let upstream = support::upstream([response.clone(), response]).await;
|
||||
let route = support::messages_route(litellm_inference_testing::no_secrets())
|
||||
.with_cache(ScopedCache::new(cache, CacheScope::Shared));
|
||||
let request = || litellm_inference_messages::MessagesCall {
|
||||
body: streaming_call.body.clone(),
|
||||
api_base: Some(upstream.uri()),
|
||||
..self::streaming_call()
|
||||
};
|
||||
let hooks = ChangingHooks::default();
|
||||
|
||||
let first = streamed_bytes(route.execute(request(), &hooks, None).await.unwrap()).await;
|
||||
let second = streamed_bytes(route.execute(request(), &hooks, None).await.unwrap()).await;
|
||||
|
||||
assert_eq!(first, wire.as_bytes());
|
||||
assert_eq!(second, wire.as_bytes());
|
||||
assert_eq!(
|
||||
support::received(&upstream).await.len(),
|
||||
if cache_hit { 1 } else { 2 }
|
||||
);
|
||||
let facts = hooks.facts.lock().unwrap();
|
||||
let [first, second] = facts.as_slice() else {
|
||||
panic!("each call reports exactly one result");
|
||||
};
|
||||
assert_eq!(first.source, ResultSource::Provider);
|
||||
assert_eq!(
|
||||
matches!(second.source, ResultSource::Cache { .. }),
|
||||
cache_hit
|
||||
);
|
||||
assert_eq!(
|
||||
second.cost,
|
||||
if cache_hit {
|
||||
Cost::CacheHit
|
||||
} else {
|
||||
Cost::Deferred
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::before_first_chunk(false)]
|
||||
#[case::after_terminal_chunk(true)]
|
||||
#[tokio::test]
|
||||
async fn dropping_a_stream_before_eof_cancels_once_and_does_not_cache(
|
||||
cache: Arc<dyn ResponseCacheService>,
|
||||
streaming_call: litellm_inference_messages::MessagesCall,
|
||||
#[case] read_chunk: bool,
|
||||
) {
|
||||
use futures_util::StreamExt;
|
||||
use litellm_host::lifecycle::{CallEvent, ExecutionEvent};
|
||||
use litellm_inference_messages::MessagesCallResponse;
|
||||
use wiremock::ResponseTemplate;
|
||||
|
||||
let response = ResponseTemplate::new(200).set_body_raw(COMPLETE_STREAM, "text/event-stream");
|
||||
let upstream = support::upstream([response.clone(), response]).await;
|
||||
let route = support::messages_route(litellm_inference_testing::no_secrets())
|
||||
.with_cache(ScopedCache::new(cache, CacheScope::Shared));
|
||||
let request = || litellm_inference_messages::MessagesCall {
|
||||
body: streaming_call.body.clone(),
|
||||
api_base: Some(upstream.uri()),
|
||||
..self::streaming_call()
|
||||
};
|
||||
let hooks = ChangingHooks::default();
|
||||
let observations = support::Observations::default();
|
||||
let first = route
|
||||
.execute(request(), &hooks, Some(observations.sender.clone()))
|
||||
.await
|
||||
.unwrap();
|
||||
let MessagesCallResponse::Stream { mut chunks, .. } = first else {
|
||||
panic!("expected a Messages stream");
|
||||
};
|
||||
if read_chunk {
|
||||
let mut delivered = Vec::new();
|
||||
while delivered.len() < COMPLETE_STREAM.len() {
|
||||
delivered.extend_from_slice(&chunks.next().await.unwrap().unwrap());
|
||||
}
|
||||
assert_eq!(delivered, COMPLETE_STREAM.as_bytes());
|
||||
}
|
||||
assert!(matches!(
|
||||
observations.lock().unwrap().as_slice(),
|
||||
[
|
||||
CallEvent::Started { .. },
|
||||
CallEvent::Execution(ExecutionEvent::ResultReady { .. })
|
||||
]
|
||||
));
|
||||
drop(chunks);
|
||||
assert!(matches!(
|
||||
observations.lock().unwrap().as_slice(),
|
||||
[
|
||||
CallEvent::Started { .. },
|
||||
CallEvent::Execution(ExecutionEvent::ResultReady { .. }),
|
||||
CallEvent::Cancelled { .. }
|
||||
]
|
||||
));
|
||||
|
||||
let second = streamed_bytes(route.execute(request(), &hooks, None).await.unwrap()).await;
|
||||
let third = streamed_bytes(route.execute(request(), &hooks, None).await.unwrap()).await;
|
||||
|
||||
assert_eq!(second, COMPLETE_STREAM.as_bytes());
|
||||
assert_eq!(third, COMPLETE_STREAM.as_bytes());
|
||||
assert_eq!(support::received(&upstream).await.len(), 2);
|
||||
let facts = hooks.facts.lock().unwrap();
|
||||
let [cancelled, completed, cached] = facts.as_slice() else {
|
||||
panic!("each call reports exactly one result");
|
||||
};
|
||||
assert_eq!(cancelled.source, ResultSource::Provider);
|
||||
assert_eq!(completed.source, ResultSource::Provider);
|
||||
assert!(matches!(cached.source, ResultSource::Cache { .. }));
|
||||
assert_eq!(cached.cost, Cost::CacheHit);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::platform("bedrock/claude_platform/claude-test")]
|
||||
#[case::mantle("bedrock_mantle/claude-test")]
|
||||
#[case::mantle_alias("bedrock/mantle/claude-test")]
|
||||
#[tokio::test]
|
||||
async fn signed_messages_bypass_the_response_cache(
|
||||
cache: Arc<dyn ResponseCacheService>,
|
||||
#[case] model: &str,
|
||||
) {
|
||||
use litellm_inference_testing::RecordingSecrets;
|
||||
use support::ReceivedRequest;
|
||||
|
||||
let response = support::json_response(json!({
|
||||
"id": "signed-message", "type": "message", "role": "assistant", "model": "claude-test",
|
||||
"content": [{"type": "text", "text": "signed reply"}], "stop_reason": "end_turn"
|
||||
}));
|
||||
let upstream = support::upstream([response.clone(), response]).await;
|
||||
let secrets = Arc::new(RecordingSecrets::new([
|
||||
("AWS_ACCESS_KEY_ID", "test-access"),
|
||||
("AWS_SECRET_ACCESS_KEY", "test-secret"),
|
||||
("AWS_SESSION_TOKEN", "test-session"),
|
||||
("AWS_REGION_NAME", "us-west-2"),
|
||||
("ANTHROPIC_AWS_WORKSPACE_ID", "test-workspace"),
|
||||
]));
|
||||
let route =
|
||||
support::messages_route(secrets).with_cache(ScopedCache::new(cache, CacheScope::Shared));
|
||||
let request = || litellm_inference_messages::MessagesCall {
|
||||
body: serde_json::from_value(json!({
|
||||
"model": model, "messages": [{"role": "user", "content": "hello"}], "max_tokens": 32
|
||||
}))
|
||||
.unwrap(),
|
||||
api_key: None,
|
||||
api_base: Some(upstream.uri()),
|
||||
..streaming_call()
|
||||
};
|
||||
let hooks = ChangingHooks::default();
|
||||
|
||||
let first = route.execute(request(), &hooks, None).await.unwrap();
|
||||
let second = route.execute(request(), &hooks, None).await.unwrap();
|
||||
|
||||
let (
|
||||
litellm_inference_messages::MessagesCallResponse::Complete(first),
|
||||
litellm_inference_messages::MessagesCallResponse::Complete(second),
|
||||
) = (first, second)
|
||||
else {
|
||||
panic!("expected completed signed messages");
|
||||
};
|
||||
assert_eq!(first, second);
|
||||
let requests = support::received(&upstream).await;
|
||||
let [first, second] = requests.as_slice() else {
|
||||
panic!("signed requests must each reach the provider");
|
||||
};
|
||||
assert!(
|
||||
first
|
||||
.header("authorization")
|
||||
.unwrap()
|
||||
.starts_with("AWS4-HMAC-SHA256 Credential=test-access/")
|
||||
);
|
||||
assert!(
|
||||
second
|
||||
.header("authorization")
|
||||
.unwrap()
|
||||
.starts_with("AWS4-HMAC-SHA256 Credential=test-access/")
|
||||
);
|
||||
let facts = hooks.facts.lock().unwrap();
|
||||
let [first, second] = facts.as_slice() else {
|
||||
panic!("each signed request reports one result");
|
||||
};
|
||||
assert_eq!(first.source, ResultSource::Provider);
|
||||
assert_eq!(second.source, ResultSource::Provider);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn oversized_response_integers_are_rejected_before_cache_storage(
|
||||
cache: Arc<dyn ResponseCacheService>,
|
||||
) {
|
||||
use litellm_inference_messages::{MessagesCall, MessagesCallResponse};
|
||||
use wiremock::ResponseTemplate;
|
||||
|
||||
let invalid = r#"{"id":"bad","type":"message","role":"assistant","model":"cache-test","content":[],"future_extension":{"wide_integer":18446744073709551617}}"#;
|
||||
let valid = json!({"id":"good","type":"message","role":"assistant","model":"cache-test","content":[],"future_extension":{"wide_integer":u64::MAX}});
|
||||
let upstream = support::upstream([
|
||||
ResponseTemplate::new(200).set_body_raw(invalid, "application/json"),
|
||||
support::json_response(valid.clone()),
|
||||
])
|
||||
.await;
|
||||
let route = support::messages_route(litellm_inference_testing::no_secrets())
|
||||
.with_cache(ScopedCache::new(cache, CacheScope::Shared));
|
||||
let request = || {
|
||||
MessagesCall {
|
||||
body: serde_json::from_value(json!({"model":"anthropic/cache-test","messages":[{"role":"user","content":"hi"}],"max_tokens":16})).unwrap(),
|
||||
api_key: Some("synthetic".into()), api_base: Some(upstream.uri()), custom_llm_provider: None,
|
||||
litellm_params: Default::default(), extra_headers: None, provider_specific_header: None,
|
||||
timeout: None, shaping: Default::default(),
|
||||
}
|
||||
};
|
||||
let rejected = route
|
||||
.execute(request(), &(), None)
|
||||
.await
|
||||
.err()
|
||||
.expect("oversized provider integer must fail");
|
||||
assert!(rejected.is_response());
|
||||
assert!(!rejected.is_request());
|
||||
let complete = route.execute(request(), &(), None).await.unwrap();
|
||||
let cached = route.execute(request(), &(), None).await.unwrap();
|
||||
let (MessagesCallResponse::Complete(complete), MessagesCallResponse::Complete(cached)) =
|
||||
(complete, cached)
|
||||
else {
|
||||
panic!("expected complete messages");
|
||||
};
|
||||
assert_eq!(
|
||||
complete.extra["future_extension"],
|
||||
valid["future_extension"]
|
||||
);
|
||||
assert_eq!(cached, complete);
|
||||
assert_eq!(support::received(&upstream).await.len(), 2);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -450,3 +450,170 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages
|
|||
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
|
||||
assert_eq!(optional_params, json!({"max_tokens": 16}));
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct GatedResult {
|
||||
entered: tokio::sync::Notify,
|
||||
release: tokio::sync::Notify,
|
||||
}
|
||||
|
||||
impl litellm_host::interceptors::Interceptors<Error> for GatedResult {
|
||||
async fn before_provider_request(
|
||||
&self,
|
||||
wire: WireRequest,
|
||||
_: RequestContext,
|
||||
) -> Result<WireRequest, Error> {
|
||||
Ok(wire)
|
||||
}
|
||||
|
||||
async fn after_provider_response(
|
||||
&self,
|
||||
_: litellm_host::interceptors::RawResponse,
|
||||
) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn result_ready(&self, _: ExecutionFacts) -> Result<(), Error> {
|
||||
self.entered.notify_one();
|
||||
self.release.notified().await;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::complete_response(false, false)]
|
||||
#[case::complete_stream(true, false)]
|
||||
#[case::cancel_response(false, true)]
|
||||
#[case::cancel_stream(true, true)]
|
||||
#[tokio::test]
|
||||
async fn pending_result_acceptance_blocks_delivery_and_cache_storage(
|
||||
call: MessagesCall,
|
||||
#[case] streaming: bool,
|
||||
#[case] cancel: bool,
|
||||
#[values(false, true)] cached: bool,
|
||||
) {
|
||||
use futures_util::TryStreamExt;
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_response::{CachePolicy, CacheScope, ResponseCache, ScopedCache};
|
||||
|
||||
let response = if streaming {
|
||||
ResponseTemplate::new(200).set_body_raw(
|
||||
"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
|
||||
"text/event-stream",
|
||||
)
|
||||
} else {
|
||||
message_response()
|
||||
};
|
||||
let upstream = upstream([response.clone(), response.clone(), response]).await;
|
||||
let route = messages_route(no_secrets()).with_cache(ScopedCache::new(
|
||||
Arc::new(ResponseCache::new(Arc::new(InMemoryCache::default()))),
|
||||
CacheScope::Shared,
|
||||
));
|
||||
let request = || {
|
||||
authenticated(
|
||||
with_fields(
|
||||
MessagesCall {
|
||||
body: call.body.clone(),
|
||||
..super::call()
|
||||
},
|
||||
json!({"stream": streaming}),
|
||||
),
|
||||
upstream.uri(),
|
||||
)
|
||||
};
|
||||
if cached {
|
||||
let output = route.execute(request(), &(), None).await.unwrap();
|
||||
if let MessagesCallResponse::Stream { chunks, .. } = output {
|
||||
chunks.try_collect::<Vec<_>>().await.unwrap();
|
||||
}
|
||||
}
|
||||
let gate = GatedResult::default();
|
||||
let observations = super::support::Observations::default();
|
||||
let mut pending = Box::pin(route.execute(request(), &gate, Some(observations.sender.clone())));
|
||||
tokio::time::timeout(Duration::from_secs(5), async {
|
||||
tokio::select! {
|
||||
_ = gate.entered.notified() => {},
|
||||
_ = &mut pending => panic!("a result cannot be delivered before acceptance"),
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(futures_util::poll!(&mut pending).is_pending());
|
||||
|
||||
let probe = RecordingHost::passthrough(request());
|
||||
let output = route
|
||||
.execute(
|
||||
request(),
|
||||
&probe,
|
||||
litellm_inference::CallOptions {
|
||||
cache: Some(CachePolicy {
|
||||
no_store: true,
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
if let MessagesCallResponse::Stream { chunks, .. } = output {
|
||||
chunks.try_collect::<Vec<_>>().await.unwrap();
|
||||
}
|
||||
assert_eq!(
|
||||
matches!(
|
||||
probe.facts.lock().unwrap()[0].source,
|
||||
ResultSource::Cache { .. }
|
||||
),
|
||||
cached
|
||||
);
|
||||
assert_eq!(received(&upstream).await.len(), if cached { 1 } else { 2 });
|
||||
|
||||
if cancel {
|
||||
drop(pending);
|
||||
let events = observations.lock().unwrap();
|
||||
assert!(matches!(events.last(), Some(CallEvent::Cancelled { .. })));
|
||||
assert_eq!(
|
||||
events
|
||||
.iter()
|
||||
.filter(|event| matches!(event, CallEvent::Cancelled { .. }))
|
||||
.count(),
|
||||
1
|
||||
);
|
||||
assert!(!events.iter().any(|event| matches!(
|
||||
event,
|
||||
CallEvent::Succeeded { .. } | CallEvent::Failed { .. }
|
||||
)));
|
||||
} else {
|
||||
gate.release.notify_one();
|
||||
let output = pending.await.unwrap();
|
||||
if let MessagesCallResponse::Stream { chunks, .. } = output {
|
||||
chunks.try_collect::<Vec<_>>().await.unwrap();
|
||||
}
|
||||
assert!(matches!(
|
||||
observations.lock().unwrap().last(),
|
||||
Some(CallEvent::Succeeded { .. })
|
||||
));
|
||||
}
|
||||
|
||||
let retry = RecordingHost::passthrough(request());
|
||||
let output = route.execute(request(), &retry, None).await.unwrap();
|
||||
if let MessagesCallResponse::Stream { chunks, .. } = output {
|
||||
chunks.try_collect::<Vec<_>>().await.unwrap();
|
||||
}
|
||||
assert_eq!(
|
||||
matches!(
|
||||
retry.facts.lock().unwrap()[0].source,
|
||||
ResultSource::Cache { .. }
|
||||
),
|
||||
cached || !cancel
|
||||
);
|
||||
assert_eq!(
|
||||
received(&upstream).await.len(),
|
||||
if cached {
|
||||
1
|
||||
} else if cancel {
|
||||
3
|
||||
} else {
|
||||
2
|
||||
}
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -477,16 +477,36 @@ async fn the_timeout_covers_a_stalled_stream_error_body(call: MessagesCall) {
|
|||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(call: MessagesCall) {
|
||||
async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(
|
||||
call: MessagesCall,
|
||||
#[values(false, true)] cached: bool,
|
||||
) {
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_response::{CacheScope, ResponseCache, ScopedCache};
|
||||
use litellm_host::lifecycle::{CallEvent, ExecutionEvent, FailureOrigin};
|
||||
|
||||
let (base, connection) = stalling_upstream().await;
|
||||
let response = messages_route(no_secrets())
|
||||
let observations = support::Observations::default();
|
||||
let route = messages_route(no_secrets());
|
||||
let route = if cached {
|
||||
route.with_cache(ScopedCache::new(
|
||||
Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
|
||||
Some(100),
|
||||
Some(Duration::from_secs(60)),
|
||||
)))),
|
||||
CacheScope::Shared,
|
||||
))
|
||||
} else {
|
||||
route
|
||||
};
|
||||
let response = route
|
||||
.execute(
|
||||
MessagesCall {
|
||||
timeout: Some(Duration::from_millis(300)),
|
||||
..streaming(call, base)
|
||||
},
|
||||
&(),
|
||||
None,
|
||||
Some(observations.sender.clone()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
|
@ -507,6 +527,18 @@ async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(call: MessagesC
|
|||
assert!(!error.is_request(), "{error:?}");
|
||||
assert_eq!(error.upstream_response(), None);
|
||||
assert!(chunks.next().await.is_none());
|
||||
drop(chunks);
|
||||
assert!(matches!(
|
||||
observations.lock().unwrap().as_slice(),
|
||||
[
|
||||
CallEvent::Started { .. },
|
||||
CallEvent::Execution(ExecutionEvent::ResultReady { .. }),
|
||||
CallEvent::Failed {
|
||||
origin: FailureOrigin::Call,
|
||||
..
|
||||
}
|
||||
]
|
||||
));
|
||||
tokio::time::timeout(Duration::from_secs(5), connection)
|
||||
.await
|
||||
.expect("the failed stream closes its upstream connection")
|
||||
|
|
|
|||
|
|
@ -67,6 +67,7 @@ impl From<ContentBlockType> for String {
|
|||
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
#[derive(Default)]
|
||||
#[serde(remote = "Self")]
|
||||
pub struct ContentBlock {
|
||||
#[serde(rename = "type", default, skip_serializing_if = "Option::is_none")]
|
||||
pub block_type: Option<ContentBlockType>,
|
||||
|
|
@ -96,6 +97,33 @@ pub struct ContentBlock {
|
|||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
impl serde::Serialize for ContentBlock {
|
||||
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
Self::serialize(self, serializer)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> serde::Deserialize<'de> for ContentBlock {
|
||||
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
|
||||
let value = Value::deserialize(deserializer)?;
|
||||
let nulls: Map<String, Value> = value
|
||||
.as_object()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|(_, value)| value.is_null())
|
||||
.map(|(key, value)| (key.clone(), value.clone()))
|
||||
.collect();
|
||||
let block = Self::deserialize(value).map_err(serde::de::Error::custom)?;
|
||||
if !matches!(block.block_type, Some(ContentBlockType::Other(_))) {
|
||||
return Ok(block);
|
||||
}
|
||||
Ok(Self {
|
||||
extra: block.extra.into_iter().chain(nulls).collect(),
|
||||
..block
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ContentBlock {
|
||||
pub fn text(text: impl Into<String>) -> Self {
|
||||
Self {
|
||||
|
|
@ -430,6 +458,12 @@ mod tests {
|
|||
}))]
|
||||
#[case::provider_specific_fields(json!({"type": "tool_use", "id": "t", "name": "f", "input": {}, "provider_specific_fields": {"x": 1}}))]
|
||||
#[case::untyped(json!({"unknown": {"nested": true}}))]
|
||||
#[case::unknown_block_with_familiar_nulls(json!({
|
||||
"type": "future_block", "text": null, "thinking": null, "signature": null,
|
||||
"data": null, "id": null, "name": null, "tool_use_id": null,
|
||||
"input": null, "content": null, "provider_specific_fields": null,
|
||||
"cache_control": null, "source": {"data": "123"}, "future": null
|
||||
}))]
|
||||
fn content_block_round_trips_unchanged(#[case] block: Value) {
|
||||
assert_eq!(round_trip::<ContentBlock>(&block), block);
|
||||
}
|
||||
|
|
@ -852,4 +886,78 @@ mod tests {
|
|||
assert_eq!(block.extra.get("source"), wire.get("source"));
|
||||
assert_eq!(serde_json::to_value(block).unwrap(), wire);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::zero(0, 0)]
|
||||
#[case::above_float_precision(9_007_199_254_740_993, 9_007_199_254_740_993)]
|
||||
#[case::integer_limits(u64::MAX, i64::MAX)]
|
||||
#[case::negative_top_k(1, i64::MIN)]
|
||||
fn request_integer_fields_preserve_exact_wire_values(
|
||||
#[case] max_tokens: u64,
|
||||
#[case] top_k: i64,
|
||||
) {
|
||||
let wire = format!(
|
||||
r#"{{"model":"test","messages":[],"max_tokens":{max_tokens},"top_k":{top_k}}}"#
|
||||
);
|
||||
let parsed: MessagesRequest = serde_json::from_str(&wire).unwrap();
|
||||
assert_eq!(parsed.params.max_tokens, Some(max_tokens));
|
||||
assert_eq!(parsed.params.top_k, Some(top_k));
|
||||
assert_eq!(
|
||||
serde_json::to_value(parsed).unwrap(),
|
||||
serde_json::from_str::<Value>(&wire).unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::negative_max_tokens(r#""max_tokens":-1"#)]
|
||||
#[case::fractional_max_tokens(r#""max_tokens":1.5"#)]
|
||||
#[case::overflow_max_tokens(r#""max_tokens":18446744073709551616"#)]
|
||||
#[case::string_max_tokens(r#""max_tokens":"16""#)]
|
||||
#[case::boolean_max_tokens(r#""max_tokens":true"#)]
|
||||
#[case::overflow_top_k(r#""top_k":9223372036854775808"#)]
|
||||
#[case::underflow_top_k(r#""top_k":-9223372036854775809"#)]
|
||||
#[case::fractional_top_k(r#""top_k":1.5"#)]
|
||||
#[case::string_stream(r#""stream":"true""#)]
|
||||
#[case::numeric_stream(r#""stream":1"#)]
|
||||
#[case::non_string_stop_sequence(r#""stop_sequences":[null]"#)]
|
||||
fn request_rejects_invalid_scalar_fields(#[case] field: &str) {
|
||||
let wire = format!(r#"{{"model":"test","messages":[],{field}}}"#);
|
||||
assert!(serde_json::from_str::<MessagesRequest>(&wire).is_err());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::missing(json!({}), json!({}))]
|
||||
#[case::null(
|
||||
json!({"max_tokens":null,"system":null,"tools":null,"stream":null,"metadata":null,"future":null}),
|
||||
json!({"future":null})
|
||||
)]
|
||||
#[case::empty(
|
||||
json!({"max_tokens":0,"system":"","tools":[],"stream":false,"metadata":{},"future":null}),
|
||||
json!({"max_tokens":0,"system":"","tools":[],"stream":false,"metadata":{},"future":null})
|
||||
)]
|
||||
fn request_distinguishes_empty_values_and_unknown_nulls_from_absent_options(
|
||||
#[case] fields: Value,
|
||||
#[case] expected: Value,
|
||||
) {
|
||||
let required = json!({"model":"test","messages":[]});
|
||||
let wire: Map<String, Value> = required
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone()
|
||||
.into_iter()
|
||||
.chain(fields.as_object().unwrap().clone())
|
||||
.collect();
|
||||
let expected: Map<String, Value> = required
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone()
|
||||
.into_iter()
|
||||
.chain(expected.as_object().unwrap().clone())
|
||||
.collect();
|
||||
let parsed: MessagesRequest = serde_json::from_value(Value::Object(wire)).unwrap();
|
||||
assert_eq!(
|
||||
serde_json::to_value(parsed).unwrap(),
|
||||
Value::Object(expected)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -169,6 +169,9 @@ mod tests {
|
|||
"type": "message_delta", "delta": {"stop_reason": "end_turn", "future_delta": 42},
|
||||
"usage": {"output_tokens": 7, "future_usage": true}
|
||||
}))]
|
||||
#[case::large_index(json!({"type":"content_block_stop","index":u64::MAX}))]
|
||||
#[case::large_usage(json!({"type":"message_stop","usage":{"input_tokens":u64::MAX,"output_tokens":9_007_199_254_740_993_u64,"future_usage":null}}))]
|
||||
#[case::partial_json(json!({"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"x\":"}}))]
|
||||
#[case::error(json!({
|
||||
"type": "error", "error": {"type": "future_error", "message": "failed", "future": 42}
|
||||
}))]
|
||||
|
|
@ -197,6 +200,14 @@ mod tests {
|
|||
#[rstest]
|
||||
#[case::invalid_index(json!({"type": "content_block_stop", "index": "zero"}))]
|
||||
#[case::missing_delta(json!({"type": "content_block_delta", "index": 0}))]
|
||||
#[case::negative_index(json!({"type":"content_block_stop","index":-1}))]
|
||||
#[case::fractional_index(json!({"type":"content_block_stop","index":0.5}))]
|
||||
#[case::null_index(json!({"type":"content_block_stop","index":null}))]
|
||||
#[case::unknown_delta(json!({"type":"content_block_delta","index":0,"delta":{"type":"future_delta","text":"x"}}))]
|
||||
#[case::tool_fragment_not_text(json!({"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":{"x":1}}}))]
|
||||
#[case::negative_usage(json!({"type":"message_stop","usage":{"output_tokens":-1}}))]
|
||||
#[case::fractional_usage(json!({"type":"message_delta","delta":{},"usage":{"input_tokens":0.5}}))]
|
||||
#[case::error_without_message(json!({"type":"error","error":{"type":"api_error"}}))]
|
||||
#[case::unknown_event(json!({"type": "future_event"}))]
|
||||
fn malformed_or_unrecognized_events_remain_rejected(#[case] wire: Value) {
|
||||
assert!(serde_json::from_value::<MessagesStreamEvent>(wire).is_err());
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ thiserror.workspace = true
|
|||
time.workspace = true
|
||||
tokio = { workspace = true, features = ["sync"] }
|
||||
url.workspace = true
|
||||
unicode-general-category = "1.1"
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
use crate::base_llm::messages::context::MessagesModelCapabilities;
|
||||
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||
use litellm_auth::CredentialPlacement;
|
||||
use litellm_secret_value::SecretValue;
|
||||
|
||||
|
|
@ -317,7 +318,12 @@ fn retain_blocks(messages: Vec<Message>, keep: impl Fn(&ContentBlock) -> bool) -
|
|||
|
||||
pub fn strip_empty_content_blocks(messages: Vec<Message>) -> Vec<Message> {
|
||||
retain_blocks(messages, |block| {
|
||||
!is_empty_text_block(block) && !is_empty_thinking_block(block)
|
||||
!is_empty_text_block(block)
|
||||
&& (!is_empty_thinking_block(block)
|
||||
|| block
|
||||
.signature
|
||||
.as_ref()
|
||||
.is_some_and(|signature| !signature.is_empty()))
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -397,6 +403,11 @@ pub fn strip_provider_specific_fields(messages: Vec<Message>) -> Vec<Message> {
|
|||
.into_iter()
|
||||
.map(|block| ContentBlock {
|
||||
provider_specific_fields: None,
|
||||
extra: block
|
||||
.extra
|
||||
.into_iter()
|
||||
.filter(|(key, _)| key != "provider_specific_fields")
|
||||
.collect(),
|
||||
..block
|
||||
})
|
||||
.collect(),
|
||||
|
|
@ -444,8 +455,45 @@ pub fn filter_billing_headers_from_system(system: SystemPrompt) -> Option<System
|
|||
}
|
||||
}
|
||||
|
||||
fn unwrapped_redacted_thinking_data(data: &str) -> Option<String> {
|
||||
let decoded = STANDARD.decode(data).ok()?;
|
||||
let native = std::str::from_utf8(&decoded).ok()?;
|
||||
if !native.is_ascii() || STANDARD.decode(native).is_err() {
|
||||
return None;
|
||||
}
|
||||
Some(native.to_string())
|
||||
}
|
||||
|
||||
fn with_native_redacted_data(block: ContentBlock) -> ContentBlock {
|
||||
if !block.is_type(ContentBlockType::RedactedThinking) {
|
||||
return block;
|
||||
}
|
||||
let Some(data) = block
|
||||
.data
|
||||
.as_deref()
|
||||
.and_then(unwrapped_redacted_thinking_data)
|
||||
else {
|
||||
return block;
|
||||
};
|
||||
ContentBlock {
|
||||
data: Some(data),
|
||||
..block
|
||||
}
|
||||
}
|
||||
|
||||
pub fn strip_encrypted_reasoning_blocks(messages: Vec<Message>) -> Vec<Message> {
|
||||
retain_blocks(messages, |block| !is_encrypted_reasoning_block(block))
|
||||
.into_iter()
|
||||
.map(|message| match message.content {
|
||||
MessageContent::Blocks(blocks) => Message {
|
||||
content: MessageContent::Blocks(
|
||||
blocks.into_iter().map(with_native_redacted_data).collect(),
|
||||
),
|
||||
..message
|
||||
},
|
||||
MessageContent::Text(_) => message,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn is_advisor_use(block: &ContentBlock) -> bool {
|
||||
|
|
@ -858,10 +906,18 @@ mod tests {
|
|||
json!([
|
||||
{"role": "user", "content": "weather?"},
|
||||
{"role": "assistant", "content": [
|
||||
{"type": "thinking", "thinking": "", "signature": "sig_abc"},
|
||||
{"type": "tool_use", "id": "toolu_01A", "name": "get_weather", "input": {"city": "Paris"}}
|
||||
]}
|
||||
])
|
||||
)]
|
||||
#[case::empty_thinking_with_empty_signature_is_removed(
|
||||
json!([{"role": "assistant", "content": [
|
||||
{"type": "thinking", "thinking": "", "signature": ""},
|
||||
{"type": "text", "text": "answer"}
|
||||
]}]),
|
||||
json!([{"role": "assistant", "content": [{"type": "text", "text": "answer"}]}])
|
||||
)]
|
||||
#[case::whitespace_thinking_beside_real_and_redacted_thinking(
|
||||
json!([{"role": "assistant", "content": [
|
||||
{"type": "thinking", "thinking": " \n "},
|
||||
|
|
@ -896,6 +952,9 @@ mod tests {
|
|||
#[case::non_empty_text(json!([{"role": "assistant", "content": [{"type": "text", "text": "hi"}]}]))]
|
||||
#[case::padded_text(json!([{"role": "assistant", "content": [{"type": "text", "text": " hi "}]}]))]
|
||||
#[case::empty_string_content(json!([{"role": "user", "content": ""}]))]
|
||||
#[case::signed_whitespace_thinking(json!([{"role": "assistant", "content": [
|
||||
{"type": "thinking", "thinking": " \n ", "signature": "sig_abc"}
|
||||
]}]))]
|
||||
#[case::textless_non_text_block(json!([{"role": "user", "content": [
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AA=="}}
|
||||
]}]))]
|
||||
|
|
@ -1001,6 +1060,10 @@ mod tests {
|
|||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", "content": "ok"}]}
|
||||
])
|
||||
)]
|
||||
#[case::unknown_block_null_provider_fields(
|
||||
json!([{"role": "user", "content": [{"type": "future", "provider_specific_fields": null, "text": null}]}]),
|
||||
json!([{"role": "user", "content": [{"type": "future", "text": null}]}])
|
||||
)]
|
||||
fn strip_provider_specific_fields_rewrites(#[case] input: Value, #[case] expected: Value) {
|
||||
assert_eq!(apply(strip_provider_specific_fields, input), expected);
|
||||
}
|
||||
|
|
@ -1075,6 +1138,33 @@ mod tests {
|
|||
assert_eq!(apply(strip_encrypted_reasoning_blocks, input), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::native(false)]
|
||||
#[case::converse_wrapped(true)]
|
||||
fn redacted_thinking_history_uses_the_native_data(#[case] wrapped: bool) {
|
||||
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||
|
||||
let native = "EmwKAhgBEgy3va3pzix/LafPsn4aDFIT2Xlxh0L5L8rLVyIwxtE3rAFBa8cr3qpP";
|
||||
let data = if wrapped {
|
||||
STANDARD.encode(native)
|
||||
} else {
|
||||
native.to_string()
|
||||
};
|
||||
let input = json!([{"role": "assistant", "content": [
|
||||
{"type": "redacted_thinking", "data": data},
|
||||
{"type": "text", "text": "answer"}
|
||||
]}]);
|
||||
assert_eq!(
|
||||
apply(strip_encrypted_reasoning_blocks, input),
|
||||
json!([
|
||||
{"role": "assistant", "content": [
|
||||
{"type": "redacted_thinking", "data": native},
|
||||
{"type": "text", "text": "answer"}
|
||||
]}
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::anthropic_signed_blocks(json!([
|
||||
{"role": "user", "content": "Solve it."},
|
||||
|
|
|
|||
|
|
@ -184,6 +184,7 @@ mod tests {
|
|||
json!([
|
||||
{"role": "user", "content": "run it"},
|
||||
{"role": "assistant", "content": [
|
||||
{"type": "thinking", "thinking": "", "signature": "sig"},
|
||||
{"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}},
|
||||
{"type": "server_tool_use", "id": "srv_1", "name": "web_search", "input": {"query": "q"}},
|
||||
{"type": "text", "text": "Web search results:\n\nURL: u"}
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ use litellm_llms_types::{
|
|||
use litellm_python_compat::{json::from_json, repr::repr, truthy::truthy};
|
||||
use serde_json::Value;
|
||||
use strum::VariantArray;
|
||||
use unicode_general_category::{GeneralCategory, get_general_category};
|
||||
|
||||
use crate::base_llm::messages::context::{
|
||||
MessagesModelCapabilities, ThinkingBudgets, ThinkingContext,
|
||||
|
|
@ -248,9 +249,34 @@ fn numeric_budget_tokens(budget: &Recognized<u64>) -> Option<i64> {
|
|||
}
|
||||
}
|
||||
|
||||
fn decimal_digit(character: char) -> Option<char> {
|
||||
if character.is_ascii_digit() {
|
||||
return Some(character);
|
||||
}
|
||||
let decimal_run_length =
|
||||
std::iter::successors(Some(u32::from(character)), |code| code.checked_sub(1))
|
||||
.take_while(|code| {
|
||||
char::from_u32(*code).is_some_and(|character| {
|
||||
get_general_category(character) == GeneralCategory::DecimalNumber
|
||||
})
|
||||
})
|
||||
.count();
|
||||
decimal_run_length
|
||||
.checked_sub(1)
|
||||
.and_then(|offset| char::from_digit((offset % 10) as u32, 10))
|
||||
}
|
||||
|
||||
/// Python's `int(str)`: surrounding whitespace, a sign, and single underscores between digits.
|
||||
fn python_int(text: &str) -> Option<i64> {
|
||||
let trimmed = text.trim();
|
||||
let normalized: String = text
|
||||
.trim()
|
||||
.chars()
|
||||
.map(|character| match character {
|
||||
'+' | '-' | '_' => Some(character),
|
||||
character => decimal_digit(character),
|
||||
})
|
||||
.collect::<Option<_>>()?;
|
||||
let trimmed = normalized.as_str();
|
||||
let digits = trimmed.trim_start_matches(['+', '-']);
|
||||
let well_formed = trimmed.len() - digits.len() <= 1
|
||||
&& !digits.is_empty()
|
||||
|
|
@ -1361,6 +1387,19 @@ mod tests {
|
|||
#[case::float_below_a_tier(json!(1023.9), "low")]
|
||||
#[case::string_budget(json!(" 16_000 "), "xhigh")]
|
||||
#[case::negative_budget(json!(-5), "low")]
|
||||
#[case::fullwidth_budget(json!("８１９２"), "xhigh")]
|
||||
#[case::arabic_indic_budget(json!("٨١٩٢"), "xhigh")]
|
||||
#[case::mixed_unicode_budget(json!(" +8_١92 "), "xhigh")]
|
||||
#[case::supplementary_decimal_budget(json!("𝟠𝟙𝟡𝟚"), "xhigh")]
|
||||
#[case::negative_unicode_budget(json!("-８_١92"), "low")]
|
||||
#[case::signed_underscored_budget(json!(" +8_192 "), "xhigh")]
|
||||
#[case::negative_underscored_budget(json!("-8_192"), "low")]
|
||||
#[case::unsigned_integer_limit(json!(u64::MAX), "xhigh")]
|
||||
#[case::positive_string_overflow(json!("184467440737095516160"), "xhigh")]
|
||||
#[case::negative_string_overflow(json!("-184467440737095516160"), "low")]
|
||||
#[case::negative_fraction_truncates_toward_zero(json!(-0.9), "low")]
|
||||
#[case::below_medium_fraction(json!(2047.9), "low")]
|
||||
#[case::at_medium_fraction(json!(2048.0), "medium")]
|
||||
fn legacy_budgets_read_like_python_numbers(#[case] budget: Value, #[case] effort: &str) {
|
||||
let translated = translate_thinking(
|
||||
request(
|
||||
|
|
@ -1380,6 +1419,16 @@ mod tests {
|
|||
|
||||
#[rstest]
|
||||
#[case::text(json!("lots"), "str")]
|
||||
#[case::superscript_digits(json!("⁸¹⁹²"), "str")]
|
||||
#[case::roman_numeral(json!("Ⅷ"), "str")]
|
||||
#[case::leading_separator(json!("_8192"), "str")]
|
||||
#[case::trailing_separator(json!("8192_"), "str")]
|
||||
#[case::repeated_separator(json!("8__192"), "str")]
|
||||
#[case::two_signs(json!("--8192"), "str")]
|
||||
#[case::sign_only(json!("+"), "str")]
|
||||
#[case::decimal_string(json!("8192.0"), "str")]
|
||||
#[case::exponent_string(json!("8e3"), "str")]
|
||||
#[case::whitespace_only(json!(" "), "str")]
|
||||
#[case::list(json!([1]), "list")]
|
||||
#[case::null(Value::Null, "NoneType")]
|
||||
#[case::boolean(json!(true), "bool")]
|
||||
|
|
|
|||
|
|
@ -388,8 +388,13 @@ fn compact_edit_from_openai(entry: &Map<String, Value>) -> Option<ContextEdit> {
|
|||
}
|
||||
let trigger = entry
|
||||
.get("compact_threshold")
|
||||
.and_then(Value::as_f64)
|
||||
.map(|threshold| json!({"type": "input_tokens", "value": threshold as i64}));
|
||||
.and_then(|threshold| match threshold {
|
||||
Value::Number(number) if number.is_i64() || number.is_u64() => Some(threshold.clone()),
|
||||
Value::Number(number) => number.as_f64().map(|number| json!(number as i64)),
|
||||
Value::Bool(value) => Some(json!(u8::from(*value))),
|
||||
_ => None,
|
||||
})
|
||||
.map(|threshold| json!({"type": "input_tokens", "value": threshold}));
|
||||
let passthrough = entry
|
||||
.iter()
|
||||
.filter(|(key, _)| !matches!(key.as_str(), "type" | "compact_threshold"))
|
||||
|
|
@ -701,6 +706,26 @@ mod tests {
|
|||
json!([{"type": "compaction", "compact_threshold": 200000}]),
|
||||
Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 200000}}]}))
|
||||
)]
|
||||
#[case::threshold_above_float_precision(
|
||||
json!([{"type": "compaction", "compact_threshold": 9_007_199_254_740_993_u64}]),
|
||||
Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 9_007_199_254_740_993_u64}}]}))
|
||||
)]
|
||||
#[case::threshold_unsigned_maximum(
|
||||
json!([{"type": "compaction", "compact_threshold": u64::MAX}]),
|
||||
Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": u64::MAX}}]}))
|
||||
)]
|
||||
#[case::threshold_signed_minimum(
|
||||
json!([{"type": "compaction", "compact_threshold": i64::MIN}]),
|
||||
Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": i64::MIN}}]}))
|
||||
)]
|
||||
#[case::true_threshold_matches_python_int(
|
||||
json!([{"type": "compaction", "compact_threshold": true}]),
|
||||
Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1}}]}))
|
||||
)]
|
||||
#[case::false_threshold_matches_python_int(
|
||||
json!([{"type": "compaction", "compact_threshold": false}]),
|
||||
Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 0}}]}))
|
||||
)]
|
||||
#[case::other_keys_pass_through(
|
||||
json!([{"type": "compaction", "compact_threshold": 150000, "instructions": "Focus on preserving code snippets"}]),
|
||||
Some(json!({"edits": [{
|
||||
|
|
|
|||
|
|
@ -198,4 +198,70 @@ mod tests {
|
|||
json!({"type": "text", "text": "reminder"})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::tool_result_first(
|
||||
json!({"role":"user","content":[{"type":"tool_result","tool_use_id":"t","content":"ok"},{"type":"text","text":"next"}],"future":null}),
|
||||
1
|
||||
)]
|
||||
#[case::tool_result_after_text(
|
||||
json!({"role":"user","content":[{"type":"text","text":"first"},{"type":"tool_result","tool_use_id":"t","content":"ok"}]}),
|
||||
2
|
||||
)]
|
||||
#[case::non_user_tool_result(
|
||||
json!({"role":"assistant","content":[{"type":"tool_result","tool_use_id":"t","content":"ok"}]}),
|
||||
2
|
||||
)]
|
||||
#[case::empty_blocks(json!({"role":"user","content":[]}), 2)]
|
||||
fn only_a_user_turn_starting_with_tool_results_moves_before_the_system_run(
|
||||
#[case] follower: Value,
|
||||
#[case] follower_index: usize,
|
||||
) {
|
||||
let first = json!({"role":"assistant","content":[{"type":"tool_use","id":"t","name":"lookup","input":{}}]});
|
||||
let converted = convert_mid_conversation_system_turns(messages(json!([
|
||||
first,
|
||||
{"role":"system","content":[{"type":"future","payload":{"keep":null}}]},
|
||||
follower
|
||||
])));
|
||||
let wire = serde_json::to_value(&converted).unwrap();
|
||||
assert_eq!(wire[0], first);
|
||||
assert_eq!(wire[follower_index], follower);
|
||||
assert_eq!(
|
||||
wire[3 - follower_index],
|
||||
json!({"role":"user","content":[
|
||||
{"type":"text","text":CONVERTED_SYSTEM_NOTE},
|
||||
{"type":"future","payload":{"keep":null}}
|
||||
]})
|
||||
);
|
||||
assert_eq!(
|
||||
convert_mid_conversation_system_turns(converted.clone()),
|
||||
converted
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn consecutive_system_runs_keep_their_order_across_multiple_tool_turns() {
|
||||
let first = json!({"role":"assistant","content":"first"});
|
||||
let result = json!({"role":"user","content":[{"type":"tool_result","tool_use_id":"t","content":"ok"}]});
|
||||
let next = json!({"role":"assistant","content":"next"});
|
||||
let converted = convert_mid_conversation_system_turns(messages(json!([
|
||||
first,
|
||||
{"role":"system","content":"one"},
|
||||
{"role":"system","content":"two"},
|
||||
result,
|
||||
next,
|
||||
{"role":"system","content":"tail"}
|
||||
])));
|
||||
assert_eq!(
|
||||
serde_json::to_value(converted).unwrap(),
|
||||
json!([
|
||||
first,
|
||||
result,
|
||||
{"role":"user","content":[{"type":"text","text":CONVERTED_SYSTEM_NOTE},{"type":"text","text":"one"}]},
|
||||
{"role":"user","content":[{"type":"text","text":CONVERTED_SYSTEM_NOTE},{"type":"text","text":"two"}]},
|
||||
next,
|
||||
{"role":"user","content":[{"type":"text","text":CONVERTED_SYSTEM_NOTE},{"type":"text","text":"tail"}]}
|
||||
])
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -160,4 +160,93 @@ mod tests {
|
|||
normalized
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::blocks(
|
||||
json!([{"type":"text","text":"context","cache_control":{"type":"ephemeral","scope":"global","ttl":"1h","future":null}}]),
|
||||
json!([{"type":"text","text":"context","cache_control":{"type":"ephemeral","ttl":"1h","future":null}}])
|
||||
)]
|
||||
#[case::text(json!("context"), json!("context"))]
|
||||
#[case::absent(Value::Null, Value::Null)]
|
||||
fn scope_removal_preserves_cache_hints_and_opaque_nested_fields(
|
||||
#[case] system: Value,
|
||||
#[case] expected_system: Value,
|
||||
) {
|
||||
let nested = json!([{"type":"text","text":"nested","cache_control":{"scope":"global"}}]);
|
||||
let tool = json!({"name":"lookup","input_schema":{},"cache_control":{"scope":"global","ttl":"1h"}});
|
||||
let request: MessagesRequest = serde_json::from_value(json!({
|
||||
"model":"test", "messages":[
|
||||
{"role":"user","content":"plain","future_message":null},
|
||||
{"role":"assistant","content":[
|
||||
{"type":"future","cache_control":{"type":"ephemeral","scope":"global","ttl":"1h","future":null},"payload":7},
|
||||
{"type":"text","text":"uncached"}
|
||||
]},
|
||||
{"role":"user","content":[{"type":"tool_result","tool_use_id":"t","content":nested,
|
||||
"cache_control":{"scope":"global","ttl":"5m"}}]}
|
||||
],
|
||||
"system":system, "tools":[tool], "future_request":null
|
||||
})).unwrap();
|
||||
let stripped = super::strip_cache_control_scope(request);
|
||||
let wire = serde_json::to_value(&stripped).unwrap();
|
||||
assert_eq!(
|
||||
wire["messages"],
|
||||
json!([
|
||||
{"role":"user","content":"plain","future_message":null},
|
||||
{"role":"assistant","content":[
|
||||
{"type":"future","cache_control":{"type":"ephemeral","ttl":"1h","future":null},"payload":7},
|
||||
{"type":"text","text":"uncached"}
|
||||
]},
|
||||
{"role":"user","content":[{"type":"tool_result","tool_use_id":"t","content":nested,
|
||||
"cache_control":{"ttl":"5m"}}]}
|
||||
])
|
||||
);
|
||||
assert_eq!(
|
||||
wire.get("system"),
|
||||
(!expected_system.is_null()).then_some(&expected_system)
|
||||
);
|
||||
assert_eq!(wire["tools"], json!([tool]));
|
||||
assert_eq!(wire.get("future_request"), Some(&Value::Null));
|
||||
assert_eq!(super::strip_cache_control_scope(stripped.clone()), stripped);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::native_mid_conversation(true, json!({"role":"system","content":"later","name":"operator"}))]
|
||||
#[case::converted_mid_conversation(false, json!({"role":"user","content":[
|
||||
{"type":"text","text":crate::base_llm::messages::mid_conversation_system::CONVERTED_SYSTEM_NOTE},
|
||||
{"type":"text","text":"later"}
|
||||
]}))]
|
||||
fn leading_system_hoisting_and_later_system_conversion_compose(
|
||||
#[case] supports_mid_conversation: bool,
|
||||
#[case] expected_later: Value,
|
||||
) {
|
||||
let request: MessagesRequest = serde_json::from_value(json!({
|
||||
"model":"test", "max_tokens":16, "system":"existing",
|
||||
"messages":[
|
||||
{"role":"system","content":"leading"},
|
||||
{"role":"assistant","content":[{"type":"tool_use","id":"t","name":"lookup","input":{}}]},
|
||||
{"role":"system","content":"later","name":"operator"},
|
||||
{"role":"user","content":[{"type":"tool_result","tool_use_id":"t","content":"ok"}]}
|
||||
]
|
||||
})).unwrap();
|
||||
let assistant = request.messages[1].clone();
|
||||
let result = request.messages[3].clone();
|
||||
let normalized = normalize_system_role_messages(request, supports_mid_conversation);
|
||||
let expected_messages = if supports_mid_conversation {
|
||||
json!([assistant, expected_later, result])
|
||||
} else {
|
||||
json!([assistant, result, expected_later])
|
||||
};
|
||||
assert_eq!(
|
||||
serde_json::to_value(&normalized).unwrap(),
|
||||
json!({
|
||||
"model":"test", "max_tokens":16,
|
||||
"system":[{"type":"text","text":"existing"},{"type":"text","text":"leading"}],
|
||||
"messages":expected_messages
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
normalize_system_role_messages(normalized.clone(), supports_mid_conversation),
|
||||
normalized
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -225,7 +225,7 @@ fn filter_body(body: Value, strip_connection_fields: bool) -> Value {
|
|||
|| name.as_str() == UNSUPPORTED_OVERRIDE))
|
||||
})
|
||||
.map(|(name, value)| {
|
||||
let value = if name == "model" {
|
||||
let value = if strip_connection_fields && name == "model" {
|
||||
value
|
||||
.as_str()
|
||||
.map(|model| {
|
||||
|
|
@ -385,6 +385,26 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::platform_route("claude_platform/claude-test")]
|
||||
#[case::bare_model("claude-test")]
|
||||
fn platform_billing_identity_is_retained_until_wire_serialization(#[case] model: &str) {
|
||||
let request = serde_json::from_value(json!({
|
||||
"model": model, "max_tokens": 16, "messages": [],
|
||||
}))
|
||||
.unwrap();
|
||||
let transformed = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG
|
||||
.transform_anthropic_messages_request(request, &MessagesTransformContext::default())
|
||||
.unwrap();
|
||||
assert_eq!(transformed.model, model);
|
||||
let wire = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG
|
||||
.wire_body(serde_json::to_value(transformed).unwrap());
|
||||
assert_eq!(
|
||||
wire,
|
||||
json!({"model":"claude-test", "max_tokens":16, "messages":[]})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::default(json!({}), false)]
|
||||
#[case::override_empty(json!({"claude_platform_unsupported_params": []}), true)]
|
||||
|
|
@ -421,4 +441,70 @@ mod tests {
|
|||
assert_eq!(wire.get("aws_region_name"), None);
|
||||
assert_eq!(wire.get(UNSUPPORTED_OVERRIDE), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::request(json!({"workspace_id": "request"}), json!({"workspace_id": "deployment"}), "request")]
|
||||
#[case::deployment(json!({}), json!({"workspace_id": "deployment"}), "deployment")]
|
||||
#[case::blank_request(json!({"workspace_id": " "}), json!({"workspace_id": "deployment"}), "deployment")]
|
||||
#[case::environment(json!({}), json!({}), "environment")]
|
||||
fn workspace_sources_select_one_header_without_changing_unrelated_headers(
|
||||
#[case] optional: Value,
|
||||
#[case] deployment: Value,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let optional_params = serde_json::from_value(optional).unwrap();
|
||||
let params = LitellmParams {
|
||||
extra: serde_json::from_value(deployment).unwrap(),
|
||||
..Default::default()
|
||||
};
|
||||
let environment = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG
|
||||
.validate_environment(
|
||||
vec![("x-routing".into(), "keep".into())],
|
||||
Some("key"),
|
||||
"claude-test",
|
||||
&optional_params,
|
||||
¶ms,
|
||||
&|name| (name == WORKSPACE_ENVS[0]).then(|| "environment".into()),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
environment.headers,
|
||||
vec![
|
||||
("x-routing".into(), "keep".into()),
|
||||
(WORKSPACE_HEADER.into(), expected.into()),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn custom_unsupported_fields_do_not_remove_nested_tool_data_or_extensions() {
|
||||
let input = json!({
|
||||
"model": "claude_platform/claude-test", "max_tokens": 16, "messages": [],
|
||||
"workspace_id": "workspace", "aws_region_name": "us-west-2",
|
||||
"claude_platform_unsupported_params": ["temperature"],
|
||||
"temperature": 0.5,
|
||||
"context_management": {"edits": [{"type": "clear_thinking_20251015"}]},
|
||||
"tools": [{"name": "lookup", "input_schema": {
|
||||
"type": "object", "properties": {"workspace_id": {"type": "string"}},
|
||||
"aws_region_name": "application-data", "temperature": 7
|
||||
}}],
|
||||
"future": null
|
||||
});
|
||||
let transformed = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG
|
||||
.transform_anthropic_messages_request(
|
||||
serde_json::from_value(input.clone()).unwrap(),
|
||||
&MessagesTransformContext::default(),
|
||||
)
|
||||
.unwrap();
|
||||
let wire = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG
|
||||
.wire_body(serde_json::to_value(transformed).unwrap());
|
||||
assert_eq!(
|
||||
wire,
|
||||
json!({
|
||||
"model": "claude-test", "max_tokens": 16, "messages": [],
|
||||
"context_management": input["context_management"],
|
||||
"tools": input["tools"], "future": null
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -690,4 +690,59 @@ mod tests {
|
|||
(other, _) => panic!("unexpected auth {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::zero_stop_counts_override_nonzero_delta(
|
||||
json!({"cache_read_input_tokens": 90, "cache_creation_input_tokens": 80}),
|
||||
json!({"input_tokens": 7, "output_tokens": 5, "cache_read_input_tokens": 4, "cache_creation_input_tokens": 3}),
|
||||
json!({"input_tokens": 0, "output_tokens": 99, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0}),
|
||||
json!({"input_tokens": 0, "output_tokens": 5, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0})
|
||||
)]
|
||||
#[case::delta_extensions_outrank_start_details(
|
||||
json!({"cache_creation": {"ephemeral_5m_input_tokens": 8}, "start_only": 1}),
|
||||
json!({"output_tokens": 5, "cache_creation": {"ephemeral_1h_input_tokens": 3}, "future": null}),
|
||||
json!({"output_tokens": 99, "stop_only": 2}),
|
||||
json!({"output_tokens": 5, "cache_creation": {"ephemeral_1h_input_tokens": 3}, "future": null})
|
||||
)]
|
||||
fn usage_precedence_preserves_explicit_zero_and_delta_extensions(
|
||||
#[case] start: Value,
|
||||
#[case] delta: Value,
|
||||
#[case] stop: Value,
|
||||
#[case] expected: Value,
|
||||
) {
|
||||
let output = promote(vec![
|
||||
message_start(start.clone()),
|
||||
message_delta(delta),
|
||||
message_stop(Some(stop.clone())),
|
||||
]);
|
||||
assert_eq!(
|
||||
output,
|
||||
vec![
|
||||
message_start(start),
|
||||
message_delta(expected),
|
||||
message_stop(Some(stop)),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn only_the_last_pending_delta_receives_terminal_usage() {
|
||||
let start = message_start(json!({"cache_read_input_tokens": 80}));
|
||||
let first = message_delta(json!({"output_tokens": 1}));
|
||||
let last = message_delta(json!({"output_tokens": 5}));
|
||||
let stop = message_stop(Some(
|
||||
json!({"input_tokens": 3, "cache_read_input_tokens": 100}),
|
||||
));
|
||||
assert_eq!(
|
||||
promote(vec![start.clone(), first.clone(), last, stop.clone()]),
|
||||
vec![
|
||||
start,
|
||||
first,
|
||||
message_delta(
|
||||
json!({"output_tokens": 5, "input_tokens": 3, "cache_read_input_tokens": 100})
|
||||
),
|
||||
stop,
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -13,7 +13,10 @@ use crate::{
|
|||
Error,
|
||||
anthropic::{
|
||||
common_utils::{filter_billing_headers_from_system, has_anthropic_credential},
|
||||
messages::transformation::update_headers_with_anthropic_beta,
|
||||
messages::{
|
||||
handler::shape_anthropic_messages_request,
|
||||
transformation::update_headers_with_anthropic_beta,
|
||||
},
|
||||
},
|
||||
base_llm::{
|
||||
auth::AuthScheme,
|
||||
|
|
@ -33,6 +36,14 @@ pub struct EdenAIAnthropicMessagesConfig;
|
|||
pub const EDENAI_MESSAGES_CONFIG: EdenAIAnthropicMessagesConfig = EdenAIAnthropicMessagesConfig;
|
||||
|
||||
impl BaseMessagesConfig for EdenAIAnthropicMessagesConfig {
|
||||
fn shape_request(
|
||||
&self,
|
||||
request: MessagesRequest,
|
||||
reasoning_auto_summary: bool,
|
||||
) -> Result<MessagesRequest, Error> {
|
||||
shape_anthropic_messages_request(request, reasoning_auto_summary)
|
||||
}
|
||||
|
||||
fn validate_environment(
|
||||
&self,
|
||||
headers: Headers,
|
||||
|
|
@ -121,3 +132,118 @@ impl BaseMessagesConfig for EdenAIAnthropicMessagesConfig {
|
|||
serde_json::Number::from_f64(cost)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
#[rstest]
|
||||
#[case::bearer("Authorization", "Bearer caller")]
|
||||
#[case::api_key("X-Api-Key", "caller")]
|
||||
fn forwarded_credentials_still_require_an_edenai_key(#[case] name: &str, #[case] value: &str) {
|
||||
let error = EDENAI_MESSAGES_CONFIG
|
||||
.validate_environment(
|
||||
vec![(name.into(), value.into())],
|
||||
None,
|
||||
"test-model",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
&|_| None,
|
||||
)
|
||||
.unwrap_err();
|
||||
assert_eq!(
|
||||
error,
|
||||
Error::Auth(litellm_auth::Error::MissingApiKey {
|
||||
provider: "Eden AI",
|
||||
environment_variable: API_KEY_ENV,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::explicit(Some("explicit"), "explicit")]
|
||||
#[case::missing(None, "environment")]
|
||||
#[case::blank(Some(" "), "environment")]
|
||||
fn configured_credentials_use_bearer_auth(
|
||||
#[case] api_key: Option<&str>,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let headers = vec![("x-routing".into(), "keep".into())];
|
||||
let environment = EDENAI_MESSAGES_CONFIG
|
||||
.validate_environment(
|
||||
headers.clone(),
|
||||
api_key,
|
||||
"test-model",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
&|name| (name == API_KEY_ENV).then(|| "environment".into()),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(environment.headers, headers);
|
||||
assert!(matches!(
|
||||
environment.auth,
|
||||
AuthScheme::Credential { placement: CredentialPlacement::Bearer, secret }
|
||||
if secret.expose() == expected
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::punctuation("functions.lookup:0", "functions_lookup_0")]
|
||||
#[case::thought_signature("lookup__thought__signature", "lookup")]
|
||||
fn shaping_normalizes_matching_tool_ids_and_preserves_native_parameters(
|
||||
#[case] source_id: &str,
|
||||
#[case] expected_id: &str,
|
||||
) {
|
||||
let input = json!({"id": source_id, "cache_control": {"scope": "application"}});
|
||||
let thinking = json!({"type": "enabled", "budget_tokens": 512, "native_field": true});
|
||||
let body = json!({
|
||||
"model": "test-model", "max_tokens": 64, "thinking": thinking,
|
||||
"messages": [
|
||||
{"role": "assistant", "content": [
|
||||
{"type": "text", "text": " \n "},
|
||||
{"type": "tool_use", "id": source_id, "name": "lookup", "input": input}
|
||||
]},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": source_id, "content": "ok"}]}
|
||||
]
|
||||
});
|
||||
let shaped = EDENAI_MESSAGES_CONFIG
|
||||
.shape_request(serde_json::from_value(body).unwrap(), false)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
serde_json::to_value(shaped).unwrap(),
|
||||
json!({
|
||||
"model": "test-model", "max_tokens": 64, "thinking": thinking,
|
||||
"messages": [
|
||||
{"role": "assistant", "content": [{"type": "tool_use", "id": expected_id, "name": "lookup", "input": input}]},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": expected_id, "content": "ok"}]}
|
||||
]
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn cache_sanitization_preserves_tool_application_data_and_extensions() {
|
||||
let payload = json!({"cache_control": {"type": "application", "ttl": "keep"}});
|
||||
let wire = EDENAI_MESSAGES_CONFIG.wire_body(json!({
|
||||
"model": "test-model",
|
||||
"messages": [{"role": "assistant", "content": [{
|
||||
"type": "tool_use", "id": "call-1", "name": "lookup", "input": payload,
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h", "scope": "global"}
|
||||
}]}],
|
||||
"future": null
|
||||
}));
|
||||
assert_eq!(
|
||||
wire,
|
||||
json!({
|
||||
"model": "test-model",
|
||||
"messages": [{"role": "assistant", "content": [{
|
||||
"type": "tool_use", "id": "call-1", "name": "lookup", "input": payload,
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
}]}],
|
||||
"future": null
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -198,4 +198,65 @@ mod tests {
|
|||
assert_eq!(body["provider"], input["provider"]);
|
||||
assert_eq!(body["thinking"], input["thinking"]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty_bearer("Bearer ")]
|
||||
#[case::basic_auth("Basic legacy")]
|
||||
#[case::unrelated_scheme("Token legacy")]
|
||||
fn an_invalid_forwarded_credential_is_replaced_without_leaking_an_api_key(
|
||||
#[case] authorization: &str,
|
||||
) {
|
||||
let environment = OPENROUTER_MESSAGES_CONFIG
|
||||
.validate_environment(
|
||||
vec![
|
||||
("AuThOrIzAtIoN".into(), authorization.into()),
|
||||
("X-Api-Key".into(), "wrong-provider-key".into()),
|
||||
("x-routing".into(), "keep".into()),
|
||||
],
|
||||
Some("deployment-key"),
|
||||
"test-model",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
&|name| (name == API_KEY_ENV).then(|| "environment-key".into()),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
environment.headers,
|
||||
vec![("x-routing".into(), "keep".into())]
|
||||
);
|
||||
assert!(matches!(
|
||||
environment.auth,
|
||||
AuthScheme::Credential { placement: CredentialPlacement::Bearer, secret }
|
||||
if secret.expose() == "deployment-key"
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::canonical("Bearer caller")]
|
||||
#[case::case_and_whitespace(" bEaReR caller ")]
|
||||
fn a_forwarded_bearer_wins_but_foreign_api_keys_are_removed(#[case] bearer: &str) {
|
||||
let environment = OPENROUTER_MESSAGES_CONFIG
|
||||
.validate_environment(
|
||||
vec![
|
||||
("Authorization".into(), bearer.into()),
|
||||
("x-api-key".into(), "foreign-key".into()),
|
||||
("X-API-KEY".into(), "another-foreign-key".into()),
|
||||
("x-routing".into(), "keep".into()),
|
||||
],
|
||||
Some("deployment-key"),
|
||||
"test-model",
|
||||
&MessagesOptionalParams::default(),
|
||||
&LitellmParams::default(),
|
||||
&|name| (name == API_KEY_ENV).then(|| "environment-key".into()),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
environment.headers,
|
||||
vec![
|
||||
("Authorization".into(), bearer.into()),
|
||||
("x-routing".into(), "keep".into()),
|
||||
]
|
||||
);
|
||||
assert!(matches!(environment.auth, AuthScheme::Forwarded));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,9 +2,10 @@ use litellm_auth::CredentialPlacement;
|
|||
use litellm_secret_value::SecretValue;
|
||||
|
||||
use litellm_core_utils::settings::resolve_non_empty;
|
||||
use litellm_http::request::without_headers;
|
||||
use litellm_llms_types::{
|
||||
formats::messages::{MessagesOptionalParams, MessagesRequest},
|
||||
providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH},
|
||||
providers::anthropic::{BETA_HEADER, DEFAULT_HEADERS, MESSAGES_PATH},
|
||||
};
|
||||
use litellm_router_types::LitellmParams;
|
||||
|
||||
|
|
@ -13,8 +14,7 @@ use crate::{
|
|||
anthropic::{
|
||||
common_utils::{filter_billing_headers_from_system, has_anthropic_credential},
|
||||
messages::{
|
||||
handler::shape_anthropic_messages_request,
|
||||
transformation::{transform_messages_request, update_headers_with_anthropic_beta},
|
||||
handler::shape_anthropic_messages_request, transformation::transform_messages_request,
|
||||
},
|
||||
},
|
||||
base_llm::{
|
||||
|
|
@ -117,8 +117,8 @@ impl BaseMessagesConfig for TencentAnthropicMessagesConfig {
|
|||
DEFAULT_HEADERS
|
||||
}
|
||||
|
||||
fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers {
|
||||
update_headers_with_anthropic_beta(headers, request)
|
||||
fn request_headers(&self, headers: Headers, _request: &MessagesRequest) -> Headers {
|
||||
without_headers(headers, &[BETA_HEADER])
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -132,6 +132,31 @@ mod tests {
|
|||
None
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::lowercase("anthropic-beta")]
|
||||
#[case::mixed_case("Anthropic-Beta")]
|
||||
fn request_headers_remove_caller_and_generated_anthropic_betas(#[case] name: &str) {
|
||||
let request = serde_json::from_value(serde_json::json!({
|
||||
"model": "test-model", "max_tokens": 64,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"output_config": {"format": {"type": "json_schema", "schema": {"type": "object"}}},
|
||||
"tools": [{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}]
|
||||
}))
|
||||
.unwrap();
|
||||
let retained = vec![("x-routing".into(), "keep".into())];
|
||||
let headers = vec![
|
||||
(
|
||||
name.into(),
|
||||
"caller-beta,structured-outputs-2025-11-13".into(),
|
||||
),
|
||||
("x-routing".into(), "keep".into()),
|
||||
];
|
||||
assert_eq!(
|
||||
TENCENT_MESSAGES_CONFIG.request_headers(headers, &request),
|
||||
retained
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::root("https://proxy.test", "https://proxy.test/v1/messages")]
|
||||
#[case::root_slash("https://proxy.test/", "https://proxy.test/v1/messages")]
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ pyo3::create_exception!(
|
|||
);
|
||||
|
||||
const MESSAGES_REQUEST_ERROR: &str = "messages_request_error";
|
||||
const MESSAGES_RESPONSE_ERROR: &str = "messages_response_error";
|
||||
/// The rejection is a parameter the model does not support, which Python raises as
|
||||
/// `UnsupportedParamsError`.
|
||||
const MESSAGES_UNSUPPORTED_PARAM: &str = "messages_unsupported_param";
|
||||
|
|
@ -137,6 +138,9 @@ pub(crate) fn route_failure(
|
|||
if error.is_missing_credential() {
|
||||
value.setattr(MESSAGES_MISSING_CREDENTIAL, true)?;
|
||||
}
|
||||
if error.is_response() {
|
||||
value.setattr(MESSAGES_RESPONSE_ERROR, true)?;
|
||||
}
|
||||
native
|
||||
}
|
||||
(RequestMarker::Unmarked, _) => by_fault(error.is_request(), &error),
|
||||
|
|
|
|||
|
|
@ -186,13 +186,11 @@ impl MessagesPythonHost {
|
|||
}
|
||||
|
||||
fn provider(&self, py: Python<'_>) -> String {
|
||||
self.view
|
||||
.request(py)
|
||||
.and_then(|request| request.get_item("custom_llm_provider"))
|
||||
.ok()
|
||||
.flatten()
|
||||
.and_then(|value| value.extract::<Option<String>>().ok().flatten())
|
||||
.unwrap_or_else(|| "anthropic".into())
|
||||
py.import(ROUTE_HOST_MODULE)
|
||||
.and_then(|module| module.getattr("request_provider"))
|
||||
.and_then(|resolve| resolve.call1((self.view.request(py)?,)))
|
||||
.and_then(|value| value.extract())
|
||||
.unwrap_or_else(|_| "anthropic".into())
|
||||
}
|
||||
|
||||
fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr {
|
||||
|
|
|
|||
|
|
@ -7,6 +7,21 @@ use pyo3::prelude::*;
|
|||
use super::NativeCall;
|
||||
|
||||
fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
|
||||
let request = call.resolved()?;
|
||||
let model: String = request
|
||||
.get_item("model")?
|
||||
.ok_or_else(|| pyo3::exceptions::PyTypeError::new_err("messages requires a model"))?
|
||||
.extract()?;
|
||||
let provider: String = py
|
||||
.import("litellm.rust_bridge.host.routes.messages")?
|
||||
.getattr("request_provider")?
|
||||
.call1((&request,))?
|
||||
.extract()?;
|
||||
if !litellm_inference_messages::supports_messages_provider(&model, Some(&provider)) {
|
||||
return Err(crate::errors::RustBridgeDeclined::new_err(format!(
|
||||
"messages provider {provider} has no native adapter"
|
||||
)));
|
||||
}
|
||||
let (arguments, hooks) =
|
||||
crate::routes::call_hooks(py, LoggingOperation::Messages, &call, asynchronous)?;
|
||||
crate::routes::run_public_call(
|
||||
|
|
|
|||
|
|
@ -69,8 +69,7 @@ def _public_request(
|
|||
return None
|
||||
model: Final = fields.get("model")
|
||||
messages: Final = optional_sequence(fields.get("messages"))
|
||||
max_tokens: Final = fields.get("max_tokens")
|
||||
if not isinstance(model, str) or messages is None or not isinstance(max_tokens, int):
|
||||
if not isinstance(model, str) or messages is None:
|
||||
return None
|
||||
return native_call(legacy, args, _native_kwargs(fields, kwargs))
|
||||
|
||||
|
|
|
|||
|
|
@ -1111,6 +1111,8 @@ class FallbackAwareAnthropicMessagesStream:
|
|||
getattr(source_iterator, "_hidden_params", None) or {}
|
||||
)
|
||||
self._followed_source_params: object = None
|
||||
self._started: bool = False
|
||||
self._closed: bool = False
|
||||
|
||||
@property
|
||||
def has_buffered_provider_output(self) -> bool:
|
||||
|
|
@ -1154,10 +1156,18 @@ class FallbackAwareAnthropicMessagesStream:
|
|||
return self
|
||||
|
||||
async def __anext__(self) -> bytes:
|
||||
self._started = True
|
||||
return await self._async_generator.__anext__()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
from litellm.llms.anthropic.messages.native.streaming_iterator import aclose_if_supported
|
||||
|
||||
if self._closed:
|
||||
return
|
||||
await self._async_generator.aclose()
|
||||
self._closed = True
|
||||
if not self._started:
|
||||
await aclose_if_supported(self._source_iterator)
|
||||
|
||||
def merge_fallback_hidden_params(
|
||||
self, fallback_hidden_params: Mapping[str, object], fallback_headers: Mapping[str, object]
|
||||
|
|
@ -1295,15 +1305,13 @@ def wrap_anthropic_stream(
|
|||
if _anthropic_stream_forwards_ping_live(chunk, has_generated_content):
|
||||
yield chunk
|
||||
continue
|
||||
if _stream_commits_now(chunk, has_generated_content, len(buffered)):
|
||||
has_generated_content = True
|
||||
# A transport can split one SSE data line across byte chunks, so pre-content detection
|
||||
# parses the held frames plus the current chunk.
|
||||
parse_window = (
|
||||
b"".join(c for c in (*buffered, chunk) if isinstance(c, (bytes, bytearray))) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes
|
||||
if not has_generated_content and isinstance(chunk, (bytes, bytearray)) # pyright: ignore[reportUnnecessaryIsInstance] # see above
|
||||
else chunk
|
||||
)
|
||||
if _stream_commits_now(parse_window, has_generated_content, len(buffered)):
|
||||
has_generated_content = True
|
||||
error_event = parse_anthropic_error_event(parse_window)
|
||||
bridged_failure = _bridged_stream_failure(source_iterator) if error_event is not None else None
|
||||
if bridged_failure is not None:
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from litellm.llms.anthropic.messages.native.handler import (
|
|||
_with_deployment_output_effort as _deployment_output_effort, # pyright: ignore[reportPrivateUsage] # the handler's own effort rounding
|
||||
)
|
||||
from litellm.llms.anthropic.messages.utils import is_reasoning_auto_summary_enabled
|
||||
from litellm.llms.bedrock.claude_platform.common_utils import CLAUDE_PLATFORM_ON_AWS_NON_REQUEST_PARAMS
|
||||
from litellm.rust_bridge import failures
|
||||
from litellm.rust_bridge.host import callbacks_legacy_python
|
||||
from litellm.rust_bridge.public_call import known_params, optional_str
|
||||
|
|
@ -84,8 +85,20 @@ def inferred_provider(model: str, api_base: str | None) -> str | None:
|
|||
return provider
|
||||
|
||||
|
||||
def request_provider(request: Mapping[str, object]) -> str:
|
||||
return (
|
||||
optional_str(request.get("custom_llm_provider"))
|
||||
or inferred_provider(str(request["model"]), optional_str(request.get("api_base")))
|
||||
or "anthropic"
|
||||
)
|
||||
|
||||
|
||||
def body_params(request: Mapping[str, object]) -> Mapping[str, object]:
|
||||
params: Final = known_params(_BODY_FIELDS, request)
|
||||
platform: Final = request_provider(request) == "bedrock" and str(request["model"]).removeprefix(
|
||||
"bedrock/"
|
||||
).startswith("claude_platform/")
|
||||
fields: Final = _BODY_FIELDS | CLAUDE_PLATFORM_ON_AWS_NON_REQUEST_PARAMS if platform else _BODY_FIELDS
|
||||
params: Final = known_params(fields, request)
|
||||
if deployment_reasoning_effort_levels(request.get("model_info")) is None:
|
||||
return params
|
||||
return _with_deployment_effort(params, request)
|
||||
|
|
@ -236,6 +249,15 @@ def map_failure(error: Exception, request: Mapping[str, object], request_provide
|
|||
_log_resolved_call(request, request_provider)
|
||||
if getattr(error, "messages_request_error", False):
|
||||
return _rejection(error, request, request_provider)
|
||||
if getattr(error, "messages_response_error", False):
|
||||
response_error: Final = litellm.APIError(
|
||||
status_code=500,
|
||||
message=str(error),
|
||||
llm_provider=request_provider,
|
||||
model=str(request["model"]).removeprefix(f"{request_provider}/"),
|
||||
)
|
||||
response_error.__context__ = error
|
||||
return response_error
|
||||
if getattr(error, "messages_missing_credential", False):
|
||||
return litellm.AuthenticationError(
|
||||
message=_MISSING_KEY,
|
||||
|
|
|
|||
|
|
@ -769,7 +769,119 @@ FALLBACK_CASES: Final = (
|
|||
Case("stream_timeout_kwarg", _kw(stream=True, stream_timeout=30), upstream=(_stream(TEXT_STREAM),)),
|
||||
)
|
||||
|
||||
CASES: Final = (*REQUEST_CASES, *RESPONSE_CASES, *FAILURE_CASES, *STREAM_CASES, *SYNC_CASES, *FALLBACK_CASES)
|
||||
CONTRACT_REGRESSION_CASES: Final = (
|
||||
Case(
|
||||
"tencent_filters_caller_and_feature_betas",
|
||||
_kw(
|
||||
extra_headers={"Anthropic-Beta": "structured-outputs-2025-11-13,caller-beta", "x-routing": "keep"},
|
||||
output_config={"format": {"type": "json_schema", "schema": {"type": "object"}}},
|
||||
tools=[{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}],
|
||||
),
|
||||
model="tencent/claude-sonnet-4-5",
|
||||
),
|
||||
Case(
|
||||
"edenai_normalizes_replayed_tool_ids",
|
||||
_kw(
|
||||
messages=[
|
||||
{"role": "user", "content": "Look it up"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": " \n "},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "functions.lookup:0",
|
||||
"name": "lookup",
|
||||
"input": {"id": "functions.lookup:0"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "tool_result", "tool_use_id": "functions.lookup:0", "content": "ok"}],
|
||||
},
|
||||
]
|
||||
),
|
||||
model="edenai/anthropic/claude-sonnet-4-5",
|
||||
),
|
||||
Case(
|
||||
"signature_only_thinking_replay",
|
||||
_kw(
|
||||
messages=[
|
||||
{"role": "user", "content": "Look it up"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": "", "signature": "sig_abc"},
|
||||
{"type": "tool_use", "id": "t", "name": "lookup", "input": {}},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t", "content": "ok"}]},
|
||||
]
|
||||
),
|
||||
),
|
||||
Case(
|
||||
"converse_wrapped_redacted_replay",
|
||||
_kw(
|
||||
messages=[
|
||||
{"role": "user", "content": "Solve it"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "redacted_thinking",
|
||||
"data": "RW13S0FoZ0JFZ3kzdmEzcHppeC9MYWZQc240YURCSVQyWGx4aDBMNUw4ckxWeUl3eHRFM3JBRkJhOGNyM3FwUA==",
|
||||
},
|
||||
{"type": "text", "text": "answer"},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Continue"},
|
||||
]
|
||||
),
|
||||
),
|
||||
Case(
|
||||
"compaction_threshold_exact_integer",
|
||||
_kw(context_management=[{"type": "compaction", "compact_threshold": 9007199254740993}]),
|
||||
),
|
||||
Case(
|
||||
"compaction_boolean_threshold",
|
||||
_kw(context_management=[{"type": "compaction", "compact_threshold": True}]),
|
||||
),
|
||||
Case(
|
||||
"unknown_content_familiar_null_fields",
|
||||
_kw(
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "future_block",
|
||||
"text": None,
|
||||
"input": None,
|
||||
"source": {"data": "123"},
|
||||
"future": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
),
|
||||
),
|
||||
Case(
|
||||
"unicode_legacy_thinking_budget",
|
||||
_kw(thinking={"type": "enabled", "budget_tokens": " +8_١92 "}, max_tokens=16000),
|
||||
model=OPUS_ADAPTIVE,
|
||||
),
|
||||
)
|
||||
|
||||
CASES: Final = (
|
||||
*REQUEST_CASES,
|
||||
*RESPONSE_CASES,
|
||||
*FAILURE_CASES,
|
||||
*STREAM_CASES,
|
||||
*SYNC_CASES,
|
||||
*FALLBACK_CASES,
|
||||
*CONTRACT_REGRESSION_CASES,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
|
|||
105
tests/test_litellm_rust/cache/test_python_cache.py
vendored
105
tests/test_litellm_rust/cache/test_python_cache.py
vendored
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
|
|
@ -20,10 +21,10 @@ from litellm.proxy.utils import InternalUsageCache
|
|||
from litellm.redis.stores.rate_limits import RateLimitScope
|
||||
from litellm.rust_bridge.catalog import Route
|
||||
from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload
|
||||
from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE
|
||||
from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE, call
|
||||
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
||||
from tests.test_litellm_rust.support.routing import native_route
|
||||
from tests.test_litellm_rust.support.routing import assert_served_natively, native_route
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
|
|
@ -361,3 +362,103 @@ async def test_legacy_cache_keeps_public_messages_responses_compatible(
|
|||
second_payload: Final = await collect(second) if stream else payload(second)
|
||||
assert second_payload == first_payload
|
||||
assert len(recording_server.requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream"))
|
||||
async def test_cancelling_a_cache_write_does_not_store_or_log_success(
|
||||
recording_server: RecordingServer, stream: bool, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.caching.base_cache import BaseCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
|
||||
entered: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
storage: Final = InMemoryCache()
|
||||
|
||||
class GatedCache(BaseCache):
|
||||
async def async_get_cache(self, key: str, **kwargs: object) -> object:
|
||||
return storage.get_cache(key, **kwargs)
|
||||
|
||||
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
|
||||
entered.set()
|
||||
await release.wait()
|
||||
storage.set_cache(key, value, **kwargs)
|
||||
|
||||
async def async_set_cache_pipeline(
|
||||
self, cache_list: list[tuple[str, object]], ttl: float | None = None
|
||||
) -> None:
|
||||
entered.set()
|
||||
await release.wait()
|
||||
await storage.async_set_cache_pipeline(cache_list, ttl=ttl)
|
||||
|
||||
monkeypatch.setattr(litellm, "cache", Cache(_backend=GatedCache()))
|
||||
recording_server.expected_requests = 2
|
||||
recording_server.default_response = (
|
||||
ResponseSpec(body=None, events=MESSAGES_EVENTS) if stream else ResponseSpec(body=MESSAGES_RESPONSE)
|
||||
)
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
async def complete_call() -> object:
|
||||
response: Final = await call(recording_server, stream=stream, callbacks=[recorder])
|
||||
if stream:
|
||||
await collect(response)
|
||||
return response
|
||||
|
||||
with native_route(Route.MESSAGES):
|
||||
task: Final = asyncio.create_task(complete_call())
|
||||
await asyncio.wait_for(entered.wait(), timeout=5)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
release.set()
|
||||
await drain_logging()
|
||||
assert storage.cache_dict == {}
|
||||
assert len(recording_server.requests) == 1
|
||||
assert not any("success" in name or "failure" in name for name in recorder.names)
|
||||
|
||||
second: Final = await call(recording_server, stream=stream)
|
||||
assert cache_key(second) is None
|
||||
if stream:
|
||||
await collect(second)
|
||||
third: Final = await call(recording_server, stream=stream)
|
||||
assert isinstance(cache_key(third), str)
|
||||
if stream:
|
||||
await collect(third)
|
||||
|
||||
assert_served_natively(recording_server, count=2)
|
||||
assert len(recording_server.requests) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("first_stream", (False, True), ids=("response_first", "stream_first"))
|
||||
@pytest.mark.parametrize("preset", (False, True), ids=("normal_key", "preset_key"))
|
||||
async def test_cached_messages_preserve_the_requested_response_mode(
|
||||
recording_server: RecordingServer, first_stream: bool, preset: bool, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "cache", Cache())
|
||||
recording_server.expected_requests = None
|
||||
|
||||
async def request(stream: bool) -> object:
|
||||
recording_server.default_response = (
|
||||
ResponseSpec(body=None, events=MESSAGES_EVENTS) if stream else ResponseSpec(body=MESSAGES_RESPONSE)
|
||||
)
|
||||
options: Final = {"litellm_params": {"preset_cache_key": "messages-mode"}} if preset else {}
|
||||
response: Final = await call(recording_server, stream=stream, **options)
|
||||
assert isinstance(response, AsyncIterator) is stream
|
||||
if stream:
|
||||
assert await collect(response) == b"".join(recording_server.default_response.payloads())
|
||||
else:
|
||||
assert isinstance(response, dict)
|
||||
assert response["id"] == MESSAGES_RESPONSE["id"]
|
||||
return response
|
||||
|
||||
with native_route(Route.MESSAGES):
|
||||
await request(first_stream)
|
||||
second: Final = await request(not first_stream)
|
||||
assert cache_key(second) is None
|
||||
third: Final = await request(not first_stream)
|
||||
assert isinstance(cache_key(third), str)
|
||||
|
||||
assert_served_natively(recording_server, count=2)
|
||||
assert len(recording_server.requests) == 2
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import litellm
|
|||
from tests.test_litellm_rust.support.clickhouse import clickhouse_url as clickhouse_url
|
||||
from tests.test_litellm_rust.support.isolation import isolated_test_state
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service
|
||||
from tests.test_litellm_rust.support.served import NativeEntry, recording_native_entries
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(autouse=True, loop_scope="function")
|
||||
|
|
@ -15,9 +16,16 @@ async def isolate_test_state() -> AsyncIterator[None]:
|
|||
yield
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def native_entries() -> Generator[list[NativeEntry]]:
|
||||
with recording_native_entries() as entries:
|
||||
yield entries
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def recording_server() -> Generator[RecordingServer]:
|
||||
def recording_server(native_entries: list[NativeEntry]) -> Generator[RecordingServer]:
|
||||
with recording_service() as server:
|
||||
server.native_entries = native_entries
|
||||
yield server
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
import asyncio
|
||||
from collections.abc import Iterator
|
||||
import math
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
|
||||
from contextvars import ContextVar
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
from typing import Final, Protocol, cast # noqa: TID251 # narrows public exception attributes from the legacy mapper
|
||||
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
|
@ -12,7 +13,7 @@ from litellm.integrations.argilla import ArgillaLogger
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.integrations.argilla import ArgillaItem
|
||||
from litellm.types.utils import ModelResponse, StandardLoggingPayload
|
||||
from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_RESPONSE, RATE_LIMITED, call
|
||||
from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_EVENTS, MESSAGES_RESPONSE, RATE_LIMITED, call
|
||||
from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body
|
||||
from tests.test_litellm_rust.support.isolation import rebound
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
|
||||
|
|
@ -424,3 +425,167 @@ async def test_sync_call_raises_like_python_without_logging(server: RecordingSer
|
|||
|
||||
assert not server.requests
|
||||
assert not recorder.names
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("explicit_provider", (False, True))
|
||||
@pytest.mark.parametrize("status", (401, 429))
|
||||
async def test_inferred_provider_is_preserved_in_native_public_errors(
|
||||
server: RecordingServer, explicit_provider: bool, status: int
|
||||
) -> None:
|
||||
server.default_response = ResponseSpec(
|
||||
body={"error": {"type": "provider_error", "message": "synthetic refusal"}}, status=status
|
||||
)
|
||||
error_type: Final = litellm.AuthenticationError if status == 401 else litellm.RateLimitError
|
||||
with pytest.raises(error_type) as raised:
|
||||
await call(
|
||||
server, model="deepseek/deepseek-v4-flash", custom_llm_provider="deepseek" if explicit_provider else None
|
||||
)
|
||||
assert_served_natively(server)
|
||||
public_error: Final = cast(
|
||||
_ProviderError, raised.value
|
||||
) # cast-ok: the exception mapper supplies these public attributes
|
||||
assert public_error.llm_provider == "deepseek"
|
||||
assert public_error.model == "deepseek/deepseek-v4-flash"
|
||||
assert "synthetic refusal" in str(raised.value)
|
||||
|
||||
|
||||
class _ProviderError(Protocol):
|
||||
@property
|
||||
def llm_provider(self) -> str: ...
|
||||
|
||||
@property
|
||||
def model(self) -> str: ...
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("value", (2**64 - 1, -(2**63), "9" * 4096, 1e30))
|
||||
async def test_native_response_preserves_supported_numbers_and_numeric_strings(
|
||||
server: RecordingServer, value: int | str | float
|
||||
) -> None:
|
||||
server.default_response = ResponseSpec(body={**MESSAGES_RESPONSE, "future_extension": {"wide_integer": value}})
|
||||
response: Final = await call(server)
|
||||
assert_served_natively(server, response)
|
||||
body: Final = _OBJECT.validate_python(response)
|
||||
actual: Final = _OBJECT.validate_python(body["future_extension"])["wide_integer"]
|
||||
assert actual == value
|
||||
assert type(actual) is type(value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("value", (2**64 + 1, -(2**63) - 1))
|
||||
async def test_native_response_rejects_oversized_integers_instead_of_returning_rounded_success(
|
||||
server: RecordingServer, value: int
|
||||
) -> None:
|
||||
server.default_response = ResponseSpec(body={**MESSAGES_RESPONSE, "future_extension": {"wide_integer": value}})
|
||||
with pytest.raises(litellm.APIError, match="response integer outside") as raised:
|
||||
await call(server)
|
||||
assert raised.value.status_code == 500
|
||||
assert not isinstance(raised.value, litellm.APIConnectionError)
|
||||
assert str(value) not in str(raised.value)
|
||||
assert_served_natively(server)
|
||||
|
||||
|
||||
class _MessagesRouter(Protocol):
|
||||
def aanthropic_messages(self, **kwargs: object) -> Awaitable[object]: ...
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("entry", ("provider_prefix", "explicit_provider", "router_alias"))
|
||||
@pytest.mark.parametrize("streaming", (False, True), ids=("unary", "stream"))
|
||||
async def test_platform_billing_identity_survives_wire_model_rewriting(
|
||||
server: RecordingServer, monkeypatch: pytest.MonkeyPatch, entry: str, streaming: bool
|
||||
) -> None:
|
||||
model: Final = "claude-billing-identity-test"
|
||||
backend: Final = f"bedrock/claude_platform/{model}"
|
||||
input_price: Final = 0.002
|
||||
output_price: Final = 0.003
|
||||
prices: Final = cast(
|
||||
dict[str, object], getattr(litellm, "model_cost")
|
||||
) # cast-ok: model_cost is the mutable runtime registry
|
||||
monkeypatch.setitem(
|
||||
prices,
|
||||
backend,
|
||||
{
|
||||
"litellm_provider": "bedrock",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": input_price,
|
||||
"output_cost_per_token": output_price,
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 64000,
|
||||
},
|
||||
)
|
||||
provider_response: Final = {**MESSAGES_RESPONSE, "model": model}
|
||||
events: Final = (
|
||||
(
|
||||
"message_start",
|
||||
{"type": "message_start", "message": {**provider_response, "content": [], "stop_reason": None}},
|
||||
),
|
||||
*MESSAGES_EVENTS[1:],
|
||||
)
|
||||
server.default_response = ResponseSpec(body=provider_response, events=events if streaming else ())
|
||||
recorder: Final = RecordingLogger()
|
||||
metadata: Final = {"caller_marker": "cost-context", "model_group": "billing-alias"}
|
||||
|
||||
async def invoke() -> object:
|
||||
if entry == "router_alias":
|
||||
factory: Final = cast(
|
||||
Callable[..., _MessagesRouter], litellm.Router
|
||||
) # cast-ok: Router exposes its public constructor and endpoint dynamically
|
||||
router: Final = factory(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "billing-alias",
|
||||
"litellm_params": {
|
||||
"model": backend,
|
||||
"api_base": server.base_url,
|
||||
"api_key": "synthetic",
|
||||
"workspace_id": "billing-workspace",
|
||||
},
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
disable_cooldowns=True,
|
||||
)
|
||||
return await router.aanthropic_messages(
|
||||
model="billing-alias",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
max_tokens=16,
|
||||
litellm_metadata=metadata,
|
||||
stream=streaming,
|
||||
)
|
||||
return await call(
|
||||
server,
|
||||
model=f"claude_platform/{model}" if entry == "explicit_provider" else backend,
|
||||
custom_llm_provider="bedrock" if entry == "explicit_provider" else None,
|
||||
workspace_id="billing-workspace",
|
||||
litellm_metadata=metadata,
|
||||
stream=streaming,
|
||||
)
|
||||
|
||||
with rebound(litellm, "callbacks", [recorder]):
|
||||
response: Final = await invoke()
|
||||
if streaming:
|
||||
stream: Final = cast(AsyncIterator[bytes], response) # cast-ok: native Messages streams yield bytes
|
||||
chunks: Final = [chunk async for chunk in stream]
|
||||
assert b"".join(chunks) == b"".join(server.default_response.payloads())
|
||||
successes: Final = await recorder.wait_for_async("async_log_success_event")
|
||||
if streaming:
|
||||
assert_served_natively(server)
|
||||
else:
|
||||
assert_served_natively(server, response)
|
||||
assert _OBJECT.validate_python(response)["model"] == model
|
||||
usage: Final = _OBJECT.validate_python(provider_response["usage"])
|
||||
expected: Final = (
|
||||
TypeAdapter(int).validate_python(usage["input_tokens"]) * input_price
|
||||
+ TypeAdapter(int).validate_python(usage["output_tokens"]) * output_price
|
||||
)
|
||||
logged: Final = _OBJECT.validate_python(successes[0].kwargs)
|
||||
standard: Final = _OBJECT.validate_python(logged["standard_logging_object"])
|
||||
assert math.isclose(TypeAdapter(float).validate_python(logged["response_cost"]), expected)
|
||||
assert math.isclose(TypeAdapter(float).validate_python(standard["response_cost"]), expected)
|
||||
pre_call: Final = _OBJECT.validate_python(recorder.wait_for("log_pre_api_call")[0].kwargs)
|
||||
assert pre_call["model"] == f"claude_platform/{model}"
|
||||
assert _OBJECT.validate_python(server.requests[0].body)["model"] == model
|
||||
params: Final = _OBJECT.validate_python(logged["litellm_params"])
|
||||
assert _OBJECT.validate_python(params["metadata"])["caller_marker"] == "cost-context"
|
||||
|
|
|
|||
|
|
@ -22,6 +22,8 @@ from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_ME
|
|||
from litellm.rust_bridge.public_call import NativeCall
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_ADVISOR_TOOL_TYPE
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
|
||||
from tests.test_litellm_rust.messages.data import call
|
||||
from tests.test_litellm_rust.support.recording_server import NativeEntry, RecordingServer
|
||||
|
||||
MESSAGES: Final = [{"role": "user", "content": "hi"}]
|
||||
PYTHON_RULES: Final[Rules] = ()
|
||||
|
|
@ -659,3 +661,63 @@ async def test_interceptor_feature_follows_the_resolved_provider() -> None:
|
|||
rules=rules,
|
||||
)
|
||||
assert result is native_result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_messages_admission_declines_before_provider_or_hooks(server: RecordingServer) -> None:
|
||||
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
|
||||
from tests.test_litellm_rust.support.isolation import rebound
|
||||
from tests.test_litellm_rust.support.routing import routed
|
||||
|
||||
recorder: Final = RecordingLogger()
|
||||
server.expected_requests = 0
|
||||
with routed(*RUST_RULES), rebound(litellm, "callbacks", [recorder]):
|
||||
with pytest.raises(RuntimeError, match=r"declined.*no native adapter"):
|
||||
await call(
|
||||
server,
|
||||
model="openai/gpt-5",
|
||||
max_tokens=16,
|
||||
messages=MESSAGES,
|
||||
model_info={"supported_endpoints": ["/v1/messages"]},
|
||||
api_key="synthetic",
|
||||
api_base=server.base_url,
|
||||
)
|
||||
assert server.requests == []
|
||||
assert recorder.names == ()
|
||||
assert server.native_entries == [NativeEntry(server.base_url, declined=True)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_python_traffic_cannot_satisfy_native_execution_evidence(server: RecordingServer) -> None:
|
||||
from tests.test_litellm_rust.support.routing import assert_served_natively, routed
|
||||
|
||||
with routed(*PYTHON_RULES):
|
||||
await call(
|
||||
server,
|
||||
model="anthropic/claude-test",
|
||||
max_tokens=16,
|
||||
messages=MESSAGES,
|
||||
api_key="synthetic",
|
||||
api_base=server.base_url,
|
||||
)
|
||||
assert server.native_entries == []
|
||||
with pytest.raises(AssertionError, match="native execution evidence"):
|
||||
assert_served_natively(server)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_entry_async_flag_executes_the_real_native_binding(server: RecordingServer) -> None:
|
||||
from tests.test_litellm_rust.support.routing import assert_served_natively, routed
|
||||
|
||||
with routed(*RUST_RULES):
|
||||
pending: Final = litellm.anthropic_messages_handler(
|
||||
model="anthropic/claude-test",
|
||||
max_tokens=16,
|
||||
messages=MESSAGES,
|
||||
api_key="synthetic",
|
||||
api_base=server.base_url,
|
||||
is_async=True,
|
||||
)
|
||||
assert inspect.isawaitable(pending)
|
||||
response: Final = await pending
|
||||
assert_served_natively(server, response)
|
||||
|
|
|
|||
93
tests/test_litellm_rust/messages/test_lifecycle.py
Normal file
93
tests/test_litellm_rust/messages/test_lifecycle.py
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from collections.abc import Iterator
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm import Router
|
||||
from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_MODEL
|
||||
from tests.test_litellm_rust.support.recording_server import RecordedRequest, RecordingServer, ResponseSpec
|
||||
from tests.test_litellm_rust.support.routing import assert_served_natively
|
||||
|
||||
pytestmark: Final = pytest.mark.requires_rust_extension
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def unfinished_messages_peer() -> Iterator[tuple[RecordingServer, threading.Event]]:
|
||||
requests: Final[list[RecordedRequest]] = [] # mutable-ok: HTTP handler records observed requests
|
||||
disconnected: Final = threading.Event()
|
||||
prefix: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS[:1]).payloads()[0]
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def do_POST(self) -> None:
|
||||
body: Final = self.rfile.read(int(self.headers.get("content-length", "0")))
|
||||
requests.append(
|
||||
RecordedRequest(
|
||||
method="POST",
|
||||
path=self.path,
|
||||
headers={name.lower(): value for name, value in self.headers.items()},
|
||||
raw_body=body,
|
||||
body=json.loads(body),
|
||||
)
|
||||
)
|
||||
self.send_response(200)
|
||||
self.send_header("content-type", "text/event-stream")
|
||||
self.send_header("content-length", str(len(prefix) + 1024))
|
||||
self.end_headers()
|
||||
self.wfile.write(prefix)
|
||||
self.wfile.flush()
|
||||
self.connection.settimeout(5)
|
||||
try:
|
||||
if self.rfile.read(1) == b"":
|
||||
disconnected.set()
|
||||
except ConnectionResetError:
|
||||
disconnected.set()
|
||||
except TimeoutError:
|
||||
pass
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
pass
|
||||
|
||||
http: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||||
http.daemon_threads = True
|
||||
thread: Final = threading.Thread(target=http.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield RecordingServer(http, requests, [], ResponseSpec(body={})), disconnected
|
||||
finally:
|
||||
http.shutdown()
|
||||
http.server_close()
|
||||
thread.join()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("read_first", [False, True], ids=["before-first-read", "after-first-read"])
|
||||
async def test_router_closing_a_messages_stream_releases_the_unfinished_provider_connection(
|
||||
unfinished_messages_peer: tuple[RecordingServer, threading.Event], read_first: bool
|
||||
) -> None:
|
||||
server, disconnected = unfinished_messages_peer
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "messages",
|
||||
"litellm_params": {"model": MESSAGES_MODEL, "api_base": server.base_url, "api_key": "fake-key"},
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
disable_cooldowns=True,
|
||||
)
|
||||
stream: Final = await router.aanthropic_messages(
|
||||
model="messages", messages=[{"role": "user", "content": "hello"}], max_tokens=64, stream=True
|
||||
)
|
||||
if read_first:
|
||||
assert await anext(stream) == ResponseSpec(body=None, events=MESSAGES_EVENTS[:1]).payloads()[0]
|
||||
await stream.aclose()
|
||||
await stream.aclose()
|
||||
|
||||
assert await asyncio.to_thread(disconnected.wait, 2), "closing the Router stream must close upstream"
|
||||
assert_served_natively(server, stream)
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await anext(stream)
|
||||
|
|
@ -146,3 +146,45 @@ async def test_optional_positional_body_parameters_reach_the_wire(server: Record
|
|||
assert response["id"] == MESSAGES_RESPONSE["id"]
|
||||
body, _ = sent(server)
|
||||
assert (body["metadata"], body["system"], body["temperature"]) == (args[3], args[6], args[7])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"workspace_field", ("workspace_id", "aws_workspace_id", "anthropic_workspace_id", "anthropic-workspace-id")
|
||||
)
|
||||
async def test_platform_workspace_and_unsupported_override_reach_the_native_adapter(
|
||||
server: RecordingServer, monkeypatch: pytest.MonkeyPatch, workspace_field: str
|
||||
) -> None:
|
||||
monkeypatch.delenv("ANTHROPIC_AWS_WORKSPACE_ID", raising=False)
|
||||
monkeypatch.delenv("ANTHROPIC_WORKSPACE_ID", raising=False)
|
||||
response: Final = await call(
|
||||
server,
|
||||
model="bedrock/claude_platform/claude-test",
|
||||
asynchronous=True,
|
||||
**{workspace_field: "deployment-workspace"},
|
||||
claude_platform_unsupported_params=["temperature"],
|
||||
temperature=0.3,
|
||||
tools=[
|
||||
{"name": "lookup", "input_schema": {"type": "object", "properties": {"temperature": {"type": "number"}}}}
|
||||
],
|
||||
)
|
||||
assert_served_natively(server, response)
|
||||
body, headers = sent(server)
|
||||
assert headers["anthropic-workspace-id"] == "deployment-workspace"
|
||||
assert "temperature" not in body
|
||||
assert "claude_platform_unsupported_params" not in body
|
||||
assert workspace_field not in body
|
||||
assert body["tools"] == [
|
||||
{"name": "lookup", "input_schema": {"type": "object", "properties": {"temperature": {"type": "number"}}}}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("max_tokens", (1.5, True, "16", None, [], {}))
|
||||
async def test_malformed_max_tokens_is_rejected_natively_before_io(server: RecordingServer, max_tokens: object) -> None:
|
||||
server.expected_requests = 0
|
||||
with pytest.raises(
|
||||
litellm.BadRequestError, match=r"invalid Anthropic messages request|missing required field: max_tokens"
|
||||
):
|
||||
await call(server, max_tokens=max_tokens)
|
||||
assert_served_natively(server, count=0)
|
||||
|
|
|
|||
|
|
@ -56,15 +56,37 @@ async def test_sync_stream_raises_like_python_before_sending(server: RecordingSe
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_closed_early_logs_success_once(server: RecordingServer) -> None:
|
||||
server.enqueue(STREAM)
|
||||
async def test_stream_closed_early_logs_only_delivered_usage_once(server: RecordingServer) -> None:
|
||||
first_event: Final = MESSAGES_EVENTS[0]
|
||||
start_message: Final = {
|
||||
**MESSAGES_RESPONSE,
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"usage": {"input_tokens": 5, "output_tokens": 0},
|
||||
}
|
||||
events: Final = (
|
||||
(first_event[0], {"type": "message_start", "message": start_message}),
|
||||
*MESSAGES_EVENTS[1:],
|
||||
)
|
||||
response: Final = ResponseSpec(body=None, events=events, chunk_delay=0.1)
|
||||
server.enqueue(response)
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
stream: Final = await call(server, stream=True, callbacks=[recorder])
|
||||
assert isinstance(stream, AsyncIterator)
|
||||
await anext(stream)
|
||||
first: Final = await anext(stream)
|
||||
assert first == response.payloads()[0]
|
||||
assert b"message_stop" not in first
|
||||
await stream.aclose()
|
||||
await stream.aclose()
|
||||
|
||||
assert len(await recorder.wait_for_async("async_log_success_event")) == 1
|
||||
assert_served_natively(server)
|
||||
success: Final = await recorder.wait_for_async("async_log_success_event")
|
||||
assert len(success) == 1
|
||||
assert success[0].response.usage.prompt_tokens == start_message["usage"]["input_tokens"]
|
||||
assert success[0].response.usage.completion_tokens == 0
|
||||
assert not any("failure" in name for name in recorder.names)
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await anext(stream)
|
||||
await drain_logging()
|
||||
assert len(tuple(event for event in recorder.events if event.name == "async_log_success_event")) == 1
|
||||
|
|
|
|||
|
|
@ -10,6 +10,12 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|||
from typing import Final
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NativeEntry:
|
||||
api_base: str | None
|
||||
declined: bool
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecordedRequest:
|
||||
method: str
|
||||
|
|
@ -45,6 +51,7 @@ class RecordingServer:
|
|||
responses: list[ResponseSpec]
|
||||
default_response: ResponseSpec
|
||||
expected_requests: int | None = 1
|
||||
native_entries: list[NativeEntry] = field(default_factory=list)
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
|
|
|
|||
|
|
@ -1,16 +1,104 @@
|
|||
from typing import Final
|
||||
from collections.abc import Awaitable, Callable, Generator
|
||||
from contextlib import contextmanager
|
||||
from typing import Final, cast # noqa: TID251 # observation adapters preserve the native callable protocol
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict
|
||||
from tests.test_litellm_rust.support.recording_server import RecordingServer
|
||||
from litellm.rust_bridge import _native
|
||||
from litellm.rust_bridge.bindings import NativeBinding
|
||||
from litellm.rust_bridge.chat_completions.entrypoints import NATIVE_ACOMPLETION, NATIVE_COMPLETION
|
||||
from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES
|
||||
from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR
|
||||
from litellm.rust_bridge.public_call import NativeCall, optional_str
|
||||
from litellm.rust_bridge.responses.entrypoints import NATIVE_ARESPONSES, NATIVE_RESPONSES
|
||||
from tests.test_litellm_rust.support.recording_server import NativeEntry, RecordingServer
|
||||
|
||||
SyncEntry = Callable[[NativeCall], object]
|
||||
AsyncEntry = Callable[[NativeCall], Awaitable[object]]
|
||||
|
||||
|
||||
def _sync_observer(native: SyncEntry, entries: list[NativeEntry]) -> SyncEntry:
|
||||
def observed(request: NativeCall) -> object:
|
||||
try:
|
||||
result: Final = native(request)
|
||||
except _native.RustBridgeDeclined:
|
||||
entries.append(NativeEntry(optional_str(request.resolved.get("api_base")), True))
|
||||
raise
|
||||
except BaseException:
|
||||
entries.append(NativeEntry(optional_str(request.resolved.get("api_base")), False))
|
||||
raise
|
||||
entries.append(NativeEntry(optional_str(request.resolved.get("api_base")), False))
|
||||
return result
|
||||
|
||||
return observed
|
||||
|
||||
|
||||
def _async_observer(native: AsyncEntry, entries: list[NativeEntry]) -> AsyncEntry:
|
||||
async def observed(request: NativeCall) -> object:
|
||||
try:
|
||||
result: Final = await native(request)
|
||||
except _native.RustBridgeDeclined:
|
||||
entries.append(NativeEntry(optional_str(request.resolved.get("api_base")), True))
|
||||
raise
|
||||
except BaseException:
|
||||
entries.append(NativeEntry(optional_str(request.resolved.get("api_base")), False))
|
||||
raise
|
||||
entries.append(NativeEntry(optional_str(request.resolved.get("api_base")), False))
|
||||
return result
|
||||
|
||||
return observed
|
||||
|
||||
|
||||
def _observe_sync_binding(binding: NativeBinding[SyncEntry], entries: list[NativeEntry]) -> None:
|
||||
native: Final = binding.load()
|
||||
if native is not None:
|
||||
binding.override(_sync_observer(native, entries))
|
||||
|
||||
|
||||
def _observe_async_binding(binding: NativeBinding[AsyncEntry], entries: list[NativeEntry]) -> None:
|
||||
native: Final = binding.load()
|
||||
if native is not None:
|
||||
binding.override(_async_observer(native, entries))
|
||||
|
||||
|
||||
@contextmanager
|
||||
def recording_native_entries() -> Generator[list[NativeEntry]]:
|
||||
entries: Final[list[NativeEntry]] = []
|
||||
synchronous: Final = cast( # cast-ok: all sync route bindings accept NativeCall and return a public result
|
||||
tuple[NativeBinding[SyncEntry], ...],
|
||||
(NATIVE_MESSAGES, NATIVE_RESPONSES, NATIVE_COMPLETION, NATIVE_OCR),
|
||||
)
|
||||
asynchronous: Final = cast( # cast-ok: all sync route bindings accept NativeCall and return a public result
|
||||
tuple[NativeBinding[AsyncEntry], ...],
|
||||
(NATIVE_AMESSAGES, NATIVE_ARESPONSES, NATIVE_ACOMPLETION, NATIVE_AOCR),
|
||||
)
|
||||
try:
|
||||
for binding in synchronous:
|
||||
_observe_sync_binding(binding, entries)
|
||||
for async_binding in asynchronous:
|
||||
_observe_async_binding(async_binding, entries)
|
||||
yield entries
|
||||
finally:
|
||||
for binding in (*synchronous, *asynchronous):
|
||||
binding.reset()
|
||||
|
||||
|
||||
def assert_served_natively(server: RecordingServer, *responses: object, count: int = 1) -> None:
|
||||
assert len(server.requests) == (len(responses) or count)
|
||||
if not responses:
|
||||
assert not any(
|
||||
request.headers.get("user-agent", "").startswith("python-httpx") for request in server.requests
|
||||
accepted: Final = tuple(
|
||||
entry
|
||||
for entry in server.native_entries
|
||||
if not entry.declined and (entry.api_base is None or entry.api_base.startswith(server.base_url))
|
||||
)
|
||||
assert len(accepted) >= max(count, 1), "no matching native execution evidence"
|
||||
for response in responses:
|
||||
additional_headers: Final = get_hidden_params_dict(response).get("additional_headers")
|
||||
assert isinstance(additional_headers, dict)
|
||||
assert additional_headers.get("x-litellm-rust") == "true"
|
||||
_assert_response_marker(response)
|
||||
|
||||
|
||||
def _assert_response_marker(response: object) -> None:
|
||||
headers: Final = TypeAdapter(dict[str, object]).validate_python(
|
||||
get_hidden_params_dict(response).get("additional_headers")
|
||||
)
|
||||
assert headers.get("x-litellm-rust") == "true"
|
||||
|
|
|
|||
|
|
@ -245,3 +245,14 @@ def test_a_late_rejection_tells_the_logger_the_resolved_call_first() -> None:
|
|||
assert isinstance(mapped, litellm.BadRequestError)
|
||||
assert logger.model_call_details["custom_llm_provider"] == "anthropic"
|
||||
assert logger.model_call_details["optional_params"] == {"max_tokens": 8}
|
||||
|
||||
|
||||
def test_invalid_provider_response_is_an_api_error_with_its_original_cause() -> None:
|
||||
error: Final = RuntimeError("invalid messages response integer outside i64/u64 range")
|
||||
error.messages_response_error = True # pyright: ignore[reportAttributeAccessIssue] # marker supplied by the native boundary
|
||||
mapped: Final = messages_host.map_failure(error, {"model": "deepseek/model"}, "deepseek")
|
||||
assert isinstance(mapped, litellm.APIError)
|
||||
assert not isinstance(mapped, litellm.APIConnectionError)
|
||||
assert mapped.status_code == 500
|
||||
assert mapped.__context__ is error
|
||||
assert "response integer outside" in str(mapped)
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ import os
|
|||
import sys
|
||||
import threading
|
||||
import warnings
|
||||
from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from types import CodeType, FrameType, MappingProxyType, SimpleNamespace
|
||||
|
|
@ -52,6 +52,7 @@ from litellm.router_rust.invoke.streaming import (
|
|||
MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS,
|
||||
MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS,
|
||||
FallbackAwareAnthropicMessagesStream,
|
||||
StreamAttempt,
|
||||
_stream_commits_now,
|
||||
_anthropic_stream_error_is_gateway_verdict,
|
||||
_stream_fallback_error_for_raised,
|
||||
|
|
@ -20331,3 +20332,103 @@ def test_router_warns_when_a_deployment_enforces_endpoints_without_a_list(caplog
|
|||
|
||||
assert "Deployment 'embed-from-config' sets enforce_supported_endpoints without" in caplog.text
|
||||
assert "Deployment 'embed-from-api' sets enforce_supported_endpoints without" in caplog.text
|
||||
|
||||
|
||||
class _PartitionedMessagesSSE(httpx.AsyncByteStream):
|
||||
def __init__(self, chunks: tuple[bytes, ...]) -> None:
|
||||
self._chunks: Final = chunks
|
||||
self.reads = 0
|
||||
self.closes = 0
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
for chunk in self._chunks:
|
||||
self.reads += 1
|
||||
yield chunk
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closes += 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("isolate_litellm_state", "_router_httpx_transport")
|
||||
@pytest.mark.parametrize("split_at", [0, 12, 24, 50], ids=["whole", "event-name", "event-end", "data"])
|
||||
@pytest.mark.parametrize("newline", [b"\n", b"\r\n"], ids=["lf", "crlf"])
|
||||
async def test_messages_stream_preserves_generated_content_across_partitions_before_a_retryable_error(
|
||||
respx_mock: respx.MockRouter, split_at: int, newline: bytes
|
||||
) -> None:
|
||||
content: Final = _anthropic_messages_content_chunk("original").replace(b"\n", newline)
|
||||
chunks: Final = (
|
||||
_anthropic_messages_message_start_chunk().replace(b"\n", newline),
|
||||
*((content[:split_at], content[split_at:]) if split_at else (content,)),
|
||||
_anthropic_messages_overloaded_error_chunk().replace(b"\n", newline),
|
||||
)
|
||||
route: Final = respx_mock.post("https://messages-partitions.test/v1/messages").mock(
|
||||
side_effect=lambda _: httpx.Response(
|
||||
200, headers={"content-type": "text/event-stream"}, stream=_PartitionedMessagesSSE(chunks)
|
||||
)
|
||||
)
|
||||
router: Final = Router(
|
||||
model_list=[{
|
||||
"model_name": "messages",
|
||||
"litellm_params": {
|
||||
"model": "anthropic/claude-sonnet-5-5", "api_key": "fake-key",
|
||||
"api_base": "https://messages-partitions.test",
|
||||
},
|
||||
}],
|
||||
num_retries=1,
|
||||
retry_after=0,
|
||||
disable_cooldowns=True,
|
||||
)
|
||||
stream: Final = await router.aanthropic_messages(
|
||||
model="messages", messages=[{"role": "user", "content": "hello"}], max_tokens=64, stream=True
|
||||
)
|
||||
received: Final = b"".join([chunk async for chunk in stream])
|
||||
|
||||
assert received == b"".join(chunks)
|
||||
assert route.call_count == 1
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _messages_stream_attempt() -> Iterator[StreamAttempt]:
|
||||
from litellm.router_config import build_router_config
|
||||
from litellm.router_python_host import python_host
|
||||
from litellm.router_rust.core import RustCore
|
||||
|
||||
core: Final = RustCore.build(build_router_config({"num_retries": 0}), python_host)
|
||||
try:
|
||||
yield StreamAttempt(
|
||||
core=core, op="aanthropic_messages", request={}, kwargs={}, group="messages",
|
||||
order=None, depth=0, level=0, deployment_id=None, attempted_retries=0,
|
||||
)
|
||||
finally:
|
||||
core.retire()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("read_first", [False, True], ids=["unread", "started"])
|
||||
async def test_messages_stream_close_releases_its_source_once_without_reading_more(
|
||||
_messages_stream_attempt: StreamAttempt, read_first: bool
|
||||
) -> None:
|
||||
first: Final = _anthropic_messages_message_start_chunk()
|
||||
source: Final = _PartitionedMessagesSSE((first, _anthropic_messages_content_chunk("unread")))
|
||||
|
||||
async def frames() -> AsyncGenerator[bytes, None]:
|
||||
try:
|
||||
async for chunk in source:
|
||||
yield chunk
|
||||
finally:
|
||||
await source.aclose()
|
||||
|
||||
stream: Final = FallbackAwareAnthropicMessagesStream(frames(), source, _messages_stream_attempt)
|
||||
assert source.reads == source.closes == 0
|
||||
if read_first:
|
||||
assert await anext(stream) == first
|
||||
await stream.aclose()
|
||||
await stream.aclose()
|
||||
|
||||
assert source.reads == int(read_first)
|
||||
assert source.closes == 1
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await anext(stream)
|
||||
assert source.reads == int(read_first)
|
||||
assert source.closes == 1
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue