mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
feat(rust): hand upstream response headers to the native Messages stream (#43178)
* ci: drop the ocr_testing job now that tests/ocr_tests is gone Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(ocr): restore the live OCR matrix and the ocr_testing job The public litellm.ocr / aocr / Router interface is unchanged by the Rust migration, so the live provider matrix still applies. Drops the stale VCR skip list for the deleted test_rust_bridge.py. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(messages): show streamed upstream headers never reach the native stream The Python handler puts the upstream response headers on the stream's _hidden_params before the first chunk so the proxy can forward them as llm_provider-* headers. The native route drops them, and this test fails on the Rust path while passing on Python. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * feat(messages): hand upstream response headers to the native stream before its first chunk The Messages route fills MessagesStreamHead from the upstream response and yields it on Open. The Python driver converts it through the protocol host and hands it to Stream and SyncStream as their _hidden_params, so a streamed native call carries additional_headers the same way the Python handler does and the proxy can forward them as llm_provider-* headers. The relay contract lives in the core crate test, the hand-off in the host-python driver test, and the header projection in the route host test, so the recording-server test that showed the gap is dropped. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * wip --------- Co-authored-by: Yujong Lee <yujong@berri.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
1d039ed009
commit
081f73f021
20 changed files with 1274 additions and 95 deletions
|
|
@ -17,8 +17,10 @@ pub(super) async fn send(
|
|||
body: &Value,
|
||||
timeout: Option<Duration>,
|
||||
) -> Result<reqwest::Response, Error> {
|
||||
let encoded = serde_json::to_vec(body)
|
||||
.map_err(|err| Error::InvalidRequest(format!("failed to encode messages body: {err}")))?;
|
||||
let builder = headers.iter().fold(
|
||||
http_client().post(url).json(body),
|
||||
http_client().post(url).body(encoded),
|
||||
|builder, (key, value)| builder.header(key, value),
|
||||
);
|
||||
let builder = match timeout {
|
||||
|
|
|
|||
|
|
@ -6,7 +6,6 @@ use std::{
|
|||
|
||||
use bytes::Bytes;
|
||||
use litellm_auth::SecretValue;
|
||||
use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider;
|
||||
use litellm_host::{
|
||||
event::{MachineEvent, RawResponse, RequestContext, WireRequest},
|
||||
host::{Demand, Host},
|
||||
|
|
@ -22,7 +21,6 @@ use serde_json::{Map, Value};
|
|||
|
||||
use super::{
|
||||
Error,
|
||||
common_utils::messages_provider_config,
|
||||
handler::{decode_response, network, provider_error, send},
|
||||
prepare::{prepare_provider_request, resolve_provider},
|
||||
types::{MessagesRequest, MessagesShaping},
|
||||
|
|
@ -54,6 +52,11 @@ pub enum MessagesOutput {
|
|||
Streamed,
|
||||
}
|
||||
|
||||
/// The upstream response as the caller sees it at stream hand-off, before any chunk.
|
||||
pub struct MessagesStreamHead {
|
||||
pub headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
pub struct Messages;
|
||||
|
||||
impl Protocol for Messages {
|
||||
|
|
@ -62,7 +65,7 @@ impl Protocol for Messages {
|
|||
type Projection = MessagesCall;
|
||||
type Op = Infallible;
|
||||
type Chunk = Bytes;
|
||||
type StreamHead = ();
|
||||
type StreamHead = MessagesStreamHead;
|
||||
}
|
||||
|
||||
impl From<MachineFault> for Error {
|
||||
|
|
@ -77,19 +80,6 @@ impl From<MachineFault> for Error {
|
|||
pub type MessagesHost = HostChannel<Messages>;
|
||||
pub type MessagesMachine = CallMachine<Messages>;
|
||||
|
||||
/// Whether this route serves the request, decided before any callback runs so a host
|
||||
/// can still run its own path.
|
||||
pub fn supports(model: &str, custom_llm_provider: Option<&str>, stream: bool) -> bool {
|
||||
let provider = get_custom_llm_provider(model, custom_llm_provider)
|
||||
.map(|resolved| resolved.custom_llm_provider)
|
||||
.or(custom_llm_provider);
|
||||
match provider {
|
||||
Some(ANTHROPIC_MESSAGES_PROVIDER) => true,
|
||||
Some(provider) => !stream && messages_provider_config(provider).is_some(),
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// The in-process host for a request already in hand. It answers projection once and
|
||||
/// observes nothing.
|
||||
pub struct LocalMessagesHost {
|
||||
|
|
@ -152,8 +142,11 @@ async fn execute(
|
|||
model: request.model.clone(),
|
||||
custom_llm_provider: request.provider.clone(),
|
||||
optional_params: Value::Object(
|
||||
call.body
|
||||
.iter()
|
||||
request
|
||||
.body
|
||||
.as_object()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|(name, _)| !matches!(name.as_str(), "model" | "messages"))
|
||||
.map(|(name, value)| (name.clone(), value.clone()))
|
||||
.collect(),
|
||||
|
|
@ -193,7 +186,14 @@ async fn relay(
|
|||
host: &MessagesHost,
|
||||
mut response: reqwest::Response,
|
||||
) -> Result<MessagesOutput, Error> {
|
||||
if host.open(()).await? == Demand::Detached {
|
||||
let head = MessagesStreamHead {
|
||||
headers: response
|
||||
.headers()
|
||||
.iter()
|
||||
.filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_string())))
|
||||
.collect(),
|
||||
};
|
||||
if host.open(head).await? == Demand::Detached {
|
||||
return Ok(MessagesOutput::Streamed);
|
||||
}
|
||||
while let Some(chunk) = response.chunk().await.map_err(network)? {
|
||||
|
|
|
|||
210
litellm-rust/crates/core/tests/messages/host.rs
Normal file
210
litellm-rust/crates/core/tests/messages/host.rs
Normal file
|
|
@ -0,0 +1,210 @@
|
|||
use std::{convert::Infallible, sync::Mutex};
|
||||
|
||||
use litellm_core::messages::route::Messages;
|
||||
use litellm_host::{
|
||||
event::{CallEvent, MachineEvent, RequestContext, WireRequest},
|
||||
host::Host,
|
||||
};
|
||||
use litellm_llms::anthropic::common_utils::AnthropicModelCapabilities;
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
type Rewrite = Box<dyn Fn(WireRequest) -> Result<WireRequest, Error> + Send + Sync>;
|
||||
|
||||
/// Projects like `LocalMessagesHost`, answers `before_send` through `rewrite`, and keeps
|
||||
/// every event the driver emits.
|
||||
struct RecordingHost {
|
||||
call: LocalMessagesHost,
|
||||
rewrite: Rewrite,
|
||||
events: Mutex<Vec<CallEvent>>,
|
||||
optional_params: Mutex<Vec<Value>>,
|
||||
}
|
||||
|
||||
impl RecordingHost {
|
||||
fn new(call: MessagesCall, rewrite: Rewrite) -> Self {
|
||||
Self {
|
||||
call: LocalMessagesHost::new(call),
|
||||
rewrite,
|
||||
events: Mutex::new(Vec::new()),
|
||||
optional_params: Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn passthrough(call: MessagesCall) -> Self {
|
||||
Self::new(call, Box::new(Ok))
|
||||
}
|
||||
|
||||
fn raw_responses(&self) -> Vec<String> {
|
||||
self.events
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.filter_map(|event| match event {
|
||||
CallEvent::Machine(MachineEvent::ResponseReceived { raw }) => {
|
||||
Some(raw.body.clone())
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl Host<Messages> for RecordingHost {
|
||||
async fn project(&self) -> Result<MessagesCall, Error> {
|
||||
self.call.project().await
|
||||
}
|
||||
|
||||
async fn custom_op(&self, op: Infallible) -> Result<(), Error> {
|
||||
match op {}
|
||||
}
|
||||
|
||||
async fn before_send(
|
||||
&self,
|
||||
wire: WireRequest,
|
||||
context: &RequestContext,
|
||||
) -> Result<WireRequest, Error> {
|
||||
self.optional_params
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(context.optional_params.clone());
|
||||
(self.rewrite)(wire)
|
||||
}
|
||||
|
||||
async fn emit(&self, event: &CallEvent) -> Result<(), Error> {
|
||||
self.events.lock().unwrap().push(event.clone());
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_through(host: &RecordingHost) -> Result<MessagesOutput, Error> {
|
||||
litellm_host::run::run(messages_machine(Arc::new(RecordingSecrets::empty())), host).await
|
||||
}
|
||||
|
||||
fn authenticated(call: MessagesCall, api_base: String) -> MessagesCall {
|
||||
MessagesCall {
|
||||
api_key: Some("sk-ant".into()),
|
||||
api_base: Some(api_base),
|
||||
..call
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn what_before_send_returns_is_what_the_provider_receives(call: MessagesCall) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let host = RecordingHost::new(
|
||||
authenticated(call, upstream.uri()),
|
||||
Box::new(|wire| {
|
||||
let mut body = wire.body;
|
||||
body["system"] = json!("added by the host");
|
||||
Ok(WireRequest {
|
||||
headers: wire
|
||||
.headers
|
||||
.into_iter()
|
||||
.chain([("x-host".to_string(), "seen".to_string())])
|
||||
.collect(),
|
||||
body,
|
||||
..wire
|
||||
})
|
||||
}),
|
||||
);
|
||||
|
||||
run_through(&host).await.expect("messages call succeeds");
|
||||
|
||||
let request = only_request(&upstream).await;
|
||||
assert_eq!(request.json()["system"], "added by the host");
|
||||
assert_eq!(request.header("x-host"), Some("seen"));
|
||||
assert_eq!(request.header("x-api-key"), Some("sk-ant"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_before_send_failure_never_sends(call: MessagesCall) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let host = RecordingHost::new(
|
||||
authenticated(call, upstream.uri()),
|
||||
Box::new(|_| Err(Error::InvalidRequest("vetoed by the host".into()))),
|
||||
);
|
||||
|
||||
let error = run_through(&host)
|
||||
.await
|
||||
.err()
|
||||
.expect("the host failure fails the call");
|
||||
|
||||
assert_eq!(error, Error::InvalidRequest("vetoed by the host".into()));
|
||||
assert!(received(&upstream).await.is_empty());
|
||||
assert!(host.raw_responses().is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_raw_upstream_text_is_emitted_once_for_a_message(call: MessagesCall) {
|
||||
let raw = message_body();
|
||||
let upstream = upstream([json_response(raw.clone())]).await;
|
||||
let host = RecordingHost::passthrough(authenticated(call, upstream.uri()));
|
||||
|
||||
let output = run_through(&host).await.expect("messages call succeeds");
|
||||
|
||||
assert!(matches!(output, MessagesOutput::Message(_)));
|
||||
let [emitted] = <[String; 1]>::try_from(host.raw_responses())
|
||||
.unwrap_or_else(|raws| panic!("expected one raw response, got {}", raws.len()));
|
||||
assert_eq!(serde_json::from_str::<Value>(&emitted).unwrap(), raw);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::upstream_error(ResponseTemplate::new(500).set_body_string("boom"))]
|
||||
#[case::stream(ResponseTemplate::new(200).set_body_raw("event: message_stop\ndata: {}\n\n", "text/event-stream"))]
|
||||
#[tokio::test]
|
||||
async fn no_raw_response_is_emitted_for_a_stream_or_a_failure(
|
||||
call: MessagesCall,
|
||||
#[case] response: ResponseTemplate,
|
||||
) {
|
||||
let upstream = upstream([response]).await;
|
||||
let mut body = call.body.clone();
|
||||
body.insert("stream".into(), json!(true));
|
||||
let host =
|
||||
RecordingHost::passthrough(authenticated(MessagesCall { body, ..call }, upstream.uri()));
|
||||
|
||||
let _ = run_through(&host).await;
|
||||
|
||||
assert_eq!(received(&upstream).await.len(), 1);
|
||||
assert!(host.raw_responses().is_empty());
|
||||
}
|
||||
|
||||
/// Python logs `optional_params` as what it is about to send, so a dropped param must
|
||||
/// not resurface in callbacks.
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_request_context_carries_the_shaped_params_without_model_or_messages(
|
||||
call: MessagesCall,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let body: Map<String, Value> = call
|
||||
.body
|
||||
.clone()
|
||||
.into_iter()
|
||||
.chain([("temperature".to_string(), json!(0.2))])
|
||||
.collect();
|
||||
let host = RecordingHost::passthrough(authenticated(
|
||||
MessagesCall {
|
||||
body,
|
||||
shaping: MessagesShaping {
|
||||
capabilities: AnthropicModelCapabilities {
|
||||
supports_sampling_params: false,
|
||||
..AnthropicModelCapabilities::default()
|
||||
},
|
||||
drop_params: true,
|
||||
..MessagesShaping::default()
|
||||
},
|
||||
..call
|
||||
},
|
||||
upstream.uri(),
|
||||
));
|
||||
|
||||
run_through(&host).await.expect("messages call succeeds");
|
||||
|
||||
let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap())
|
||||
.unwrap_or_else(|seen| panic!("before_send runs once, saw {}", seen.len()));
|
||||
assert_eq!(optional_params, json!({"max_tokens": 16}));
|
||||
}
|
||||
|
|
@ -14,6 +14,7 @@ use wiremock::ResponseTemplate;
|
|||
mod support;
|
||||
use support::*;
|
||||
|
||||
mod host;
|
||||
mod request;
|
||||
mod response;
|
||||
mod secrets;
|
||||
|
|
|
|||
|
|
@ -1,3 +1,7 @@
|
|||
use litellm_llms::anthropic::common_utils::{
|
||||
ANTHROPIC_ADVISOR_TOOL_TYPE, ANTHROPIC_OAUTH_BETA_HEADER, AnthropicModelCapabilities,
|
||||
SupportedEffortTiers, beta,
|
||||
};
|
||||
use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders};
|
||||
use rstest::rstest;
|
||||
|
||||
|
|
@ -132,8 +136,8 @@ async fn each_provider_posts_to_its_messages_endpoint(
|
|||
assert_eq!(request.method.as_str(), "POST");
|
||||
assert_eq!(request.url.path(), path);
|
||||
assert_eq!(request.json()["model"], MODEL);
|
||||
assert_eq!(request.header("anthropic-version"), Some("2023-06-01"));
|
||||
assert_eq!(request.header("content-type"), Some("application/json"));
|
||||
assert_eq!(request.header_values("anthropic-version"), ["2023-06-01"]);
|
||||
assert_eq!(request.header_values("content-type"), ["application/json"]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
@ -249,21 +253,423 @@ async fn additional_drop_params_remove_fields_before_sending(call: MessagesCall)
|
|||
assert_eq!(sent["top_k"], 3);
|
||||
}
|
||||
|
||||
fn with_fields(call: MessagesCall, fields: Value) -> MessagesCall {
|
||||
let body: Map<String, Value> = call.body.into_iter().chain(object(fields)).collect();
|
||||
MessagesCall { body, ..call }
|
||||
}
|
||||
|
||||
fn sent_betas(request: &wiremock::Request) -> Vec<String> {
|
||||
let [header] = <[&str; 1]>::try_from(request.header_values("anthropic-beta"))
|
||||
.unwrap_or_else(|values| panic!("expected one anthropic-beta header, got {values:?}"));
|
||||
header
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.map(str::to_string)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::anthropic_streams(MODEL, Some("anthropic"), true, true)]
|
||||
#[case::anthropic_prefix_streams("anthropic/claude-sonnet-4-5", None, true, true)]
|
||||
#[case::azure_without_stream(MODEL, Some("azure_ai"), false, true)]
|
||||
#[case::azure_stream(MODEL, Some("azure_ai"), true, false)]
|
||||
#[case::other_provider(MODEL, Some("openai"), false, false)]
|
||||
#[case::unresolvable_model("no-such-model", None, false, false)]
|
||||
fn supports_matches_what_the_route_can_serve(
|
||||
#[case] model: &str,
|
||||
#[case] provider: Option<&str>,
|
||||
#[case] stream: bool,
|
||||
#[case] supported: bool,
|
||||
#[case::structured_output(json!({"output_format": {"type": "json_schema"}}), &[beta::STRUCTURED_OUTPUT])]
|
||||
#[case::fast_mode(json!({"speed": "fast"}), &[beta::FAST_MODE_2026_02_01])]
|
||||
#[case::compaction(json!({"compaction": {"enabled": true}}), &[beta::COMPACT_2026_09_04])]
|
||||
#[case::context_management_edits(
|
||||
json!({"context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}}),
|
||||
&[beta::CONTEXT_MANAGEMENT_2025_06_27]
|
||||
)]
|
||||
#[case::per_message_output_config(
|
||||
json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}),
|
||||
&[beta::PER_TURN_CONTROL_2026_07_01]
|
||||
)]
|
||||
#[case::advisor_tool(
|
||||
json!({"tools": [{"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", "model": MODEL}]}),
|
||||
&[beta::ADVISOR_TOOL_2026_03_01]
|
||||
)]
|
||||
#[case::several_features_at_once(
|
||||
json!({"speed": "fast", "output_format": {"type": "json_schema"}}),
|
||||
&[beta::STRUCTURED_OUTPUT, beta::FAST_MODE_2026_02_01]
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn feature_betas_join_the_callers_betas_in_one_sorted_header(
|
||||
call: MessagesCall,
|
||||
#[case] fields: Value,
|
||||
#[case] features: &[&str],
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let capabilities = AnthropicModelCapabilities {
|
||||
supports_speed: true,
|
||||
..AnthropicModelCapabilities::default()
|
||||
};
|
||||
|
||||
run_message(with_fields(
|
||||
MessagesCall {
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]),
|
||||
shaping: MessagesShaping {
|
||||
capabilities,
|
||||
..MessagesShaping::default()
|
||||
},
|
||||
..call
|
||||
},
|
||||
fields,
|
||||
))
|
||||
.await;
|
||||
|
||||
let sent = sent_betas(&only_request(&upstream).await);
|
||||
let mut expected: Vec<String> = features
|
||||
.iter()
|
||||
.map(|feature| feature.to_string())
|
||||
.chain(["caller-beta-2025-01-01".to_string()])
|
||||
.collect();
|
||||
expected.sort();
|
||||
assert_eq!(sent, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: MessagesCall) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
run_message(MessagesCall {
|
||||
api_key: Some("sk-ant-oat01-token".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
})
|
||||
.await;
|
||||
|
||||
let request = only_request(&upstream).await;
|
||||
assert_eq!(
|
||||
litellm_core::messages::route::supports(model, provider, stream),
|
||||
supported
|
||||
request.header("anthropic-dangerous-direct-browser-access"),
|
||||
Some("true")
|
||||
);
|
||||
assert_eq!(sent_betas(&request), [ANTHROPIC_OAUTH_BETA_HEADER]);
|
||||
assert_eq!(request.header("x-api-key"), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::anthropic("anthropic")]
|
||||
#[case::azure_ai("azure_ai")]
|
||||
#[tokio::test]
|
||||
async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[case] provider: &str) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
run_message(MessagesCall {
|
||||
custom_llm_provider: Some(provider.into()),
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
extra_headers: headers([
|
||||
("Anthropic-Version", "2024-01-01"),
|
||||
("Content-Type", "application/json; charset=utf-8"),
|
||||
]),
|
||||
..call
|
||||
})
|
||||
.await;
|
||||
|
||||
let request = only_request(&upstream).await;
|
||||
assert_eq!(request.header_values("anthropic-version"), ["2024-01-01"]);
|
||||
assert_eq!(
|
||||
request.header_values("content-type"),
|
||||
["application/json; charset=utf-8"]
|
||||
);
|
||||
}
|
||||
|
||||
fn sampling_removed() -> AnthropicModelCapabilities {
|
||||
AnthropicModelCapabilities {
|
||||
supports_sampling_params: false,
|
||||
..AnthropicModelCapabilities::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::sampling_params(sampling_removed(), json!({"temperature": 0.2, "top_p": 0.9, "top_k": 5}), &["temperature", "top_p", "top_k"], "temperature=0.2")]
|
||||
#[case::speed(AnthropicModelCapabilities::default(), json!({"speed": "fast"}), &["speed"], "speed='fast'")]
|
||||
#[tokio::test]
|
||||
async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_it(
|
||||
call: MessagesCall,
|
||||
#[case] capabilities: AnthropicModelCapabilities,
|
||||
#[case] fields: Value,
|
||||
#[case] dropped: &[&str],
|
||||
#[case] rejected_as: &str,
|
||||
) {
|
||||
let upstream = upstream([message_response(), message_response()]).await;
|
||||
let shaped = |drop_params: bool| {
|
||||
with_fields(
|
||||
MessagesCall {
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
shaping: MessagesShaping {
|
||||
capabilities: capabilities.clone(),
|
||||
drop_params,
|
||||
..MessagesShaping::default()
|
||||
},
|
||||
body: call.body.clone(),
|
||||
custom_llm_provider: call.custom_llm_provider.clone(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
model: call.model.clone(),
|
||||
timeout: call.timeout,
|
||||
},
|
||||
fields.clone(),
|
||||
)
|
||||
};
|
||||
|
||||
let error = run(shaped(false))
|
||||
.await
|
||||
.err()
|
||||
.expect("an unsupported param is rejected without drop_params");
|
||||
assert!(
|
||||
matches!(&error, Error::InvalidRequest(message) if message.contains(rejected_as)),
|
||||
"{error:?}"
|
||||
);
|
||||
assert!(received(&upstream).await.is_empty());
|
||||
|
||||
run_message(shaped(true)).await;
|
||||
let sent = only_request(&upstream).await.json();
|
||||
for name in dropped {
|
||||
assert_eq!(sent.get(*name), None, "{name} must be dropped");
|
||||
}
|
||||
assert_eq!(sent["max_tokens"], 16);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::adaptive_thinking(json!({"type": "adaptive"}), json!({"type": "adaptive", "display": "summarized"}))]
|
||||
#[case::disabled_thinking(json!({"type": "disabled"}), json!({"type": "disabled"}))]
|
||||
#[tokio::test]
|
||||
async fn reasoning_auto_summary_marks_active_thinking_on_the_wire(
|
||||
call: MessagesCall,
|
||||
#[case] thinking: Value,
|
||||
#[case] expected: Value,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
run_message(with_fields(
|
||||
MessagesCall {
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
shaping: MessagesShaping {
|
||||
capabilities: AnthropicModelCapabilities {
|
||||
supports_reasoning: true,
|
||||
supports_adaptive_thinking: true,
|
||||
..AnthropicModelCapabilities::default()
|
||||
},
|
||||
reasoning_auto_summary: true,
|
||||
..MessagesShaping::default()
|
||||
},
|
||||
..call
|
||||
},
|
||||
json!({"thinking": thinking}),
|
||||
))
|
||||
.await;
|
||||
|
||||
assert_eq!(only_request(&upstream).await.json()["thinking"], expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::reasoning_effort_on_an_adaptive_model(
|
||||
AnthropicModelCapabilities {
|
||||
supports_reasoning: true,
|
||||
supports_adaptive_thinking: true,
|
||||
supports_output_config: true,
|
||||
effort_tiers: SupportedEffortTiers { high: true, ..SupportedEffortTiers::default() },
|
||||
..AnthropicModelCapabilities::default()
|
||||
},
|
||||
json!({"reasoning_effort": "high"}),
|
||||
json!({"thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "high"}})
|
||||
)]
|
||||
#[case::reasoning_effort_on_a_legacy_model_caps_the_budget_below_max_tokens(
|
||||
AnthropicModelCapabilities {
|
||||
supports_reasoning: true,
|
||||
..AnthropicModelCapabilities::default()
|
||||
},
|
||||
json!({"reasoning_effort": "high"}),
|
||||
json!({"thinking": {"type": "enabled", "budget_tokens": 2999}})
|
||||
)]
|
||||
#[case::adaptive_payload_on_a_legacy_model_becomes_a_capped_budget(
|
||||
AnthropicModelCapabilities {
|
||||
supports_reasoning: true,
|
||||
..AnthropicModelCapabilities::default()
|
||||
},
|
||||
json!({"thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}, "temperature": 0}),
|
||||
json!({"thinking": {"type": "enabled", "budget_tokens": 2999}})
|
||||
)]
|
||||
#[case::adaptive_payload_on_a_model_without_reasoning_is_dropped(
|
||||
AnthropicModelCapabilities::default(),
|
||||
json!({"thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}),
|
||||
json!({})
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn reasoning_is_translated_by_the_model_capabilities(
|
||||
call: MessagesCall,
|
||||
#[case] capabilities: AnthropicModelCapabilities,
|
||||
#[case] fields: Value,
|
||||
#[case] expected: Value,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
run_message(with_fields(
|
||||
MessagesCall {
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
shaping: MessagesShaping {
|
||||
capabilities,
|
||||
..MessagesShaping::default()
|
||||
},
|
||||
..call
|
||||
},
|
||||
[("max_tokens".to_string(), json!(3000))]
|
||||
.into_iter()
|
||||
.chain(object(fields))
|
||||
.collect(),
|
||||
))
|
||||
.await;
|
||||
|
||||
let sent = only_request(&upstream).await.json();
|
||||
assert_eq!(sent.get("reasoning_effort"), None);
|
||||
assert_eq!(sent.get("temperature"), None);
|
||||
let reasoning: Map<String, Value> = ["thinking", "output_config"]
|
||||
.into_iter()
|
||||
.filter_map(|name| Some((name.to_string(), sent.get(name)?.clone())))
|
||||
.collect();
|
||||
assert_eq!(Value::Object(reasoning), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty_text_blocks(
|
||||
json!([{"role": "assistant", "content": [{"type": "text", "text": " "}, {"type": "text", "text": "kept"}]}]),
|
||||
json!([{"role": "assistant", "content": [{"type": "text", "text": "kept"}]}])
|
||||
)]
|
||||
#[case::provider_specific_fields(
|
||||
json!([{"role": "assistant", "content": [{"type": "text", "text": "kept", "provider_specific_fields": {"x": 1}}]}]),
|
||||
json!([{"role": "assistant", "content": [{"type": "text", "text": "kept"}]}])
|
||||
)]
|
||||
#[case::unencrypted_web_search_results_become_text(
|
||||
json!([{"role": "assistant", "content": [{
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_1",
|
||||
"content": [{"type": "web_search_result", "title": "T", "url": "https://e.x", "page_age": null}]
|
||||
}]}]),
|
||||
json!([{"role": "assistant", "content": [{"type": "text", "text": "Web search results:\n\nTitle: T\nURL: https://e.x"}]}])
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn replayed_history_is_cleaned_before_sending(
|
||||
call: MessagesCall,
|
||||
#[case] history: Value,
|
||||
#[case] expected: Value,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
run_message(with_fields(
|
||||
MessagesCall {
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
},
|
||||
json!({"messages": history}),
|
||||
))
|
||||
.await;
|
||||
|
||||
assert_eq!(only_request(&upstream).await.json()["messages"], expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn metadata_is_reduced_to_the_user_id(call: MessagesCall) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
run_message(with_fields(
|
||||
MessagesCall {
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
},
|
||||
json!({"metadata": {"user_id": "u-1", "trace_id": "internal", "tags": ["a"]}}),
|
||||
))
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
only_request(&upstream).await.json()["metadata"],
|
||||
json!({"user_id": "u-1"})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::numeric_user_id(json!({"metadata": {"user_id": 7}}))]
|
||||
#[case::missing_max_tokens(json!({"max_tokens": null}))]
|
||||
#[tokio::test]
|
||||
async fn an_invalid_request_fails_before_sending(call: MessagesCall, #[case] fields: Value) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
let error = run(with_fields(
|
||||
MessagesCall {
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
},
|
||||
fields,
|
||||
))
|
||||
.await
|
||||
.err()
|
||||
.expect("the request is rejected");
|
||||
|
||||
assert!(error.is_request(), "{error:?}");
|
||||
assert!(received(&upstream).await.is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn azure_folds_system_role_messages_into_the_system_prompt(call: MessagesCall) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
run_message(with_fields(
|
||||
MessagesCall {
|
||||
custom_llm_provider: Some("azure_ai".into()),
|
||||
api_key: Some("sk-azure".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
},
|
||||
json!({
|
||||
"system": "top level",
|
||||
"messages": [
|
||||
{"role": "system", "content": "from a message"},
|
||||
{"role": "user", "content": "hi"}
|
||||
]
|
||||
}),
|
||||
))
|
||||
.await;
|
||||
|
||||
let sent = only_request(&upstream).await.json();
|
||||
assert_eq!(
|
||||
sent["system"],
|
||||
json!([
|
||||
{"type": "text", "text": "top level"},
|
||||
{"type": "text", "text": "from a message"}
|
||||
])
|
||||
);
|
||||
assert_eq!(sent["messages"], json!([{"role": "user", "content": "hi"}]));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::bare_model(MODEL, MODEL)]
|
||||
#[case::one_prefix("anthropic/claude-sonnet-4-5", MODEL)]
|
||||
#[case::doubled_prefix_loses_one_segment(
|
||||
"anthropic/anthropic/claude-sonnet-4-5",
|
||||
"anthropic/claude-sonnet-4-5"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn the_provider_prefix_is_stripped_exactly_once(
|
||||
call: MessagesCall,
|
||||
#[case] model: &str,
|
||||
#[case] sent_model: &str,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
run_message(MessagesCall {
|
||||
model: model.into(),
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_eq!(only_request(&upstream).await.json()["model"], sent_model);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,11 +5,14 @@ use rstest::rstest;
|
|||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::anthropic("anthropic")]
|
||||
#[case::azure_ai("azure_ai")]
|
||||
#[tokio::test]
|
||||
async fn the_provider_message_is_returned(call: MessagesCall) {
|
||||
async fn the_provider_message_is_returned(call: MessagesCall, #[case] provider: &str) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
let message = run_message(MessagesCall {
|
||||
custom_llm_provider: Some(provider.into()),
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
|
|
@ -21,6 +24,88 @@ async fn the_provider_message_is_returned(call: MessagesCall) {
|
|||
assert_eq!(message.stop_reason.as_deref(), Some("end_turn"));
|
||||
}
|
||||
|
||||
/// A refusal and fields the route does not model come back exactly as the provider sent
|
||||
/// them, since the Python side returns the raw message and the router decides what to do.
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_message_passes_through_losslessly(call: MessagesCall) {
|
||||
let upstream_body = json!({
|
||||
"id": "msg_2",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": MODEL,
|
||||
"content": [
|
||||
{"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "q"}},
|
||||
{"type": "text", "text": "no", "citations": [{"type": "web_search_result_location", "url": "https://e.x"}]}
|
||||
],
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": null,
|
||||
"stop_details": {"type": "safeguard", "safeguard_types": ["dangerous_tool_use"]},
|
||||
"container": {"id": "container_1", "expires_at": "2026-01-01T00:00:00Z"},
|
||||
"context_management": {"applied_edits": []},
|
||||
"usage": {"input_tokens": 1, "output_tokens": 2, "server_tool_use": {"web_search_requests": 1}},
|
||||
"unknown_future_field": {"nested": true}
|
||||
});
|
||||
let upstream = upstream([json_response(upstream_body.clone())]).await;
|
||||
|
||||
let message = run_message(MessagesCall {
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_eq!(message.stop_reason.as_deref(), Some("refusal"));
|
||||
assert_eq!(serde_json::to_value(&message).unwrap(), upstream_body);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_json_error_envelope_is_kept_verbatim(call: MessagesCall) {
|
||||
let envelope =
|
||||
json!({"type": "error", "error": {"type": "invalid_request_error", "message": "bad"}});
|
||||
let upstream = upstream([status_response(400, envelope.clone())]).await;
|
||||
|
||||
let error = run(MessagesCall {
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
})
|
||||
.await
|
||||
.err()
|
||||
.expect("upstream error propagates");
|
||||
|
||||
let Error::Transport(TransportError::Http { status, body }) = error else {
|
||||
panic!("{error:?}");
|
||||
};
|
||||
assert_eq!(status, 400);
|
||||
assert_eq!(serde_json::from_str::<Value>(&body).unwrap(), envelope);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_long_error_body_is_truncated_at_the_documented_cap(call: MessagesCall) {
|
||||
let long = "x".repeat(600);
|
||||
let upstream = upstream([ResponseTemplate::new(500).set_body_string(long.clone())]).await;
|
||||
|
||||
let error = run(MessagesCall {
|
||||
api_key: Some("sk".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
})
|
||||
.await
|
||||
.err()
|
||||
.expect("upstream error propagates");
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
Error::Transport(TransportError::Http {
|
||||
status: 500,
|
||||
body: format!("{}... (truncated)", &long[..256])
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::bad_request(400)]
|
||||
#[case::unauthorized(401)]
|
||||
|
|
|
|||
|
|
@ -1,23 +1,30 @@
|
|||
use litellm_llms::{
|
||||
anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
|
||||
azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG,
|
||||
base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig,
|
||||
};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::anthropic("anthropic", &ANTHROPIC_MESSAGES_CONFIG, "ANTHROPIC_API_KEY", "ANTHROPIC_BASE_URL", "/v1/messages")]
|
||||
#[case::azure_ai("azure_ai", &AZURE_ANTHROPIC_MESSAGES_CONFIG, "AZURE_API_KEY", "AZURE_API_BASE", "/anthropic/v1/messages")]
|
||||
#[case::anthropic(
|
||||
"anthropic",
|
||||
"ANTHROPIC_API_KEY",
|
||||
"ANTHROPIC_BASE_URL",
|
||||
"/v1/messages",
|
||||
&["ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"]
|
||||
)]
|
||||
#[case::azure_ai(
|
||||
"azure_ai",
|
||||
"AZURE_API_KEY",
|
||||
"AZURE_API_BASE",
|
||||
"/anthropic/v1/messages",
|
||||
&["AZURE_API_KEY", "AZURE_API_BASE"]
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn the_credential_and_base_come_from_the_secret_source(
|
||||
call: MessagesCall,
|
||||
#[case] provider: &str,
|
||||
#[case] config: &dyn BaseAnthropicMessagesConfig,
|
||||
#[case] key_name: &str,
|
||||
#[case] base_name: &str,
|
||||
#[case] path: &str,
|
||||
#[case] looked_up: &[&str],
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let base = upstream.uri();
|
||||
|
|
@ -40,7 +47,7 @@ async fn the_credential_and_base_come_from_the_secret_source(
|
|||
let request = only_request(&upstream).await;
|
||||
assert_eq!(request.url.path(), path);
|
||||
assert_eq!(request.header("x-api-key"), Some("sk-from-manager"));
|
||||
assert_eq!(secrets.requested(), config.secret_names());
|
||||
assert_eq!(secrets.requested(), looked_up);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
@ -92,3 +99,102 @@ async fn a_secret_manager_failure_fails_the_call_before_sending(call: MessagesCa
|
|||
);
|
||||
assert!(received(&upstream).await.is_empty());
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum Base {
|
||||
Upstream,
|
||||
Unreachable,
|
||||
Blank,
|
||||
Absent,
|
||||
}
|
||||
|
||||
fn base_value(base: Base, upstream: &str) -> Option<String> {
|
||||
match base {
|
||||
Base::Upstream => Some(upstream.to_string()),
|
||||
Base::Unreachable => Some(UNREACHABLE_BASE.to_string()),
|
||||
Base::Blank => Some(" ".to_string()),
|
||||
Base::Absent => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::api_base_beats_base_url(Base::Upstream, Base::Unreachable)]
|
||||
#[case::blank_api_base_falls_through_to_base_url(Base::Blank, Base::Upstream)]
|
||||
#[case::base_url_alone(Base::Absent, Base::Upstream)]
|
||||
#[tokio::test]
|
||||
async fn the_anthropic_base_env_precedence_picks_the_upstream(
|
||||
call: MessagesCall,
|
||||
#[case] api_base: Base,
|
||||
#[case] base_url: Base,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let uri = upstream.uri();
|
||||
let values: Vec<(&str, &str)> = [
|
||||
("ANTHROPIC_API_KEY", Some("sk-env".to_string())),
|
||||
("ANTHROPIC_API_BASE", base_value(api_base, &uri)),
|
||||
("ANTHROPIC_BASE_URL", base_value(base_url, &uri)),
|
||||
]
|
||||
.iter()
|
||||
.filter_map(|(name, value)| Some((*name, value.as_deref()?)))
|
||||
.map(|(name, value)| (name, Box::leak(value.to_string().into_boxed_str()) as &str))
|
||||
.collect();
|
||||
|
||||
run_with(Arc::new(RecordingSecrets::new(values)), call)
|
||||
.await
|
||||
.expect("messages call reaches the upstream the precedence picks");
|
||||
|
||||
assert_eq!(only_request(&upstream).await.url.path(), "/v1/messages");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::auth_token_alone(
|
||||
&[("ANTHROPIC_AUTH_TOKEN", "tok")],
|
||||
("authorization", "Bearer tok"),
|
||||
"x-api-key"
|
||||
)]
|
||||
#[case::api_key_beats_the_auth_token(
|
||||
&[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "tok")],
|
||||
("x-api-key", "sk-env"),
|
||||
"authorization"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn the_auth_token_env_is_a_bearer_only_without_a_key(
|
||||
call: MessagesCall,
|
||||
#[case] values: &[(&str, &str)],
|
||||
#[case] expected: (&str, &str),
|
||||
#[case] absent: &str,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
run_with(
|
||||
Arc::new(RecordingSecrets::new(values.iter().copied())),
|
||||
MessagesCall {
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("messages call succeeds");
|
||||
|
||||
let request = only_request(&upstream).await;
|
||||
let (name, value) = expected;
|
||||
assert_eq!(request.header_values(name), [value]);
|
||||
assert_eq!(request.header(absent), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn azure_without_a_base_anywhere_fails_before_sending(call: MessagesCall) {
|
||||
let error = run_with(
|
||||
Arc::new(RecordingSecrets::new([("AZURE_API_KEY", "sk-azure")])),
|
||||
MessagesCall {
|
||||
custom_llm_provider: Some("azure_ai".into()),
|
||||
..call
|
||||
},
|
||||
)
|
||||
.await
|
||||
.err()
|
||||
.expect("azure needs a base");
|
||||
|
||||
assert_eq!(error, Error::Auth(litellm_auth::Error::MissingAzureApiBase));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,16 +1,25 @@
|
|||
use std::{convert::Infallible, sync::Mutex};
|
||||
|
||||
use bytes::Bytes;
|
||||
use litellm_core::messages::route::Messages;
|
||||
use litellm_core::messages::route::{Messages, MessagesStreamHead};
|
||||
use litellm_host::host::{Demand, Host};
|
||||
use rstest::rstest;
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::TcpListener,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
|
||||
const UPSTREAM_HEADERS: [(&str, &str); 2] = [
|
||||
("request-id", "req_upstream_123"),
|
||||
("anthropic-ratelimit-requests-remaining", "41"),
|
||||
];
|
||||
|
||||
const SSE_BODY: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
|
||||
|
||||
enum Seen {
|
||||
Open,
|
||||
Open(Vec<(String, String)>),
|
||||
Deliver(Bytes),
|
||||
}
|
||||
|
||||
|
|
@ -50,8 +59,8 @@ impl Host<Messages> for RecordingStreamHost {
|
|||
match op {}
|
||||
}
|
||||
|
||||
async fn open(&self, (): ()) -> Result<Demand, Error> {
|
||||
Ok(self.record(Seen::Open))
|
||||
async fn open(&self, head: MessagesStreamHead) -> Result<Demand, Error> {
|
||||
Ok(self.record(Seen::Open(head.headers)))
|
||||
}
|
||||
|
||||
async fn deliver(&self, chunk: Bytes) -> Result<Demand, Error> {
|
||||
|
|
@ -71,7 +80,10 @@ fn streaming(call: MessagesCall, api_base: String) -> MessagesCall {
|
|||
}
|
||||
|
||||
fn sse_response() -> ResponseTemplate {
|
||||
ResponseTemplate::new(200).set_body_raw(SSE_BODY, "text/event-stream")
|
||||
UPSTREAM_HEADERS.iter().fold(
|
||||
ResponseTemplate::new(200).set_body_raw(SSE_BODY, "text/event-stream"),
|
||||
|response, (name, value)| response.insert_header(*name, *value),
|
||||
)
|
||||
}
|
||||
|
||||
async fn stream_through(host: &RecordingStreamHost) -> Result<MessagesOutput, Error> {
|
||||
|
|
@ -80,7 +92,7 @@ async fn stream_through(host: &RecordingStreamHost) -> Result<MessagesOutput, Er
|
|||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_stream_opens_once_before_relaying_the_upstream_body(call: MessagesCall) {
|
||||
async fn upstream_headers_are_on_the_stream_head_before_the_first_chunk(call: MessagesCall) {
|
||||
let upstream = upstream([sse_response()]).await;
|
||||
let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX);
|
||||
|
||||
|
|
@ -88,14 +100,24 @@ async fn the_stream_opens_once_before_relaying_the_upstream_body(call: MessagesC
|
|||
|
||||
assert!(matches!(outcome, MessagesOutput::Streamed));
|
||||
let seen = host.seen.into_inner().unwrap();
|
||||
let [Seen::Open, chunks @ ..] = seen.as_slice() else {
|
||||
let [Seen::Open(headers), chunks @ ..] = seen.as_slice() else {
|
||||
panic!("the stream opens before any chunk is delivered");
|
||||
};
|
||||
let surfaced: Vec<(&str, &str)> = headers
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
UPSTREAM_HEADERS
|
||||
.iter()
|
||||
.any(|(upstream, _)| upstream == name)
|
||||
})
|
||||
.map(|(name, value)| (name.as_str(), value.as_str()))
|
||||
.collect();
|
||||
assert_eq!(surfaced, UPSTREAM_HEADERS);
|
||||
let delivered: Vec<u8> = chunks
|
||||
.iter()
|
||||
.flat_map(|step| match step {
|
||||
Seen::Deliver(chunk) => chunk.to_vec(),
|
||||
Seen::Open => panic!("the stream opens exactly once"),
|
||||
Seen::Open(_) => panic!("the stream opens exactly once"),
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(delivered, SSE_BODY.as_bytes());
|
||||
|
|
@ -118,9 +140,18 @@ async fn a_detached_caller_receives_nothing_more(call: MessagesCall, #[case] det
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::text_body(ResponseTemplate::new(429).set_body_string("slow down"), "slow down")]
|
||||
#[case::json_envelope(
|
||||
status_response(429, json!({"type": "error", "error": {"type": "rate_limit_error", "message": "slow down"}})),
|
||||
r#"{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}"#
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn an_upstream_error_fails_the_call_without_opening_the_stream(call: MessagesCall) {
|
||||
let upstream = upstream([ResponseTemplate::new(429).set_body_string("slow down")]).await;
|
||||
async fn an_upstream_error_fails_the_call_without_opening_the_stream(
|
||||
call: MessagesCall,
|
||||
#[case] response: ResponseTemplate,
|
||||
#[case] body: &str,
|
||||
) {
|
||||
let upstream = upstream([response]).await;
|
||||
let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX);
|
||||
|
||||
let error = stream_through(&host)
|
||||
|
|
@ -128,16 +159,88 @@ async fn an_upstream_error_fails_the_call_without_opening_the_stream(call: Messa
|
|||
.err()
|
||||
.expect("upstream error propagates");
|
||||
|
||||
assert!(
|
||||
matches!(
|
||||
error,
|
||||
Error::Transport(litellm_http::transport::Error::Http { status: 429, .. })
|
||||
),
|
||||
"{error:?}"
|
||||
assert_eq!(
|
||||
error,
|
||||
Error::Transport(litellm_http::transport::Error::Http {
|
||||
status: 429,
|
||||
body: body.into()
|
||||
})
|
||||
);
|
||||
assert!(host.seen.into_inner().unwrap().is_empty());
|
||||
}
|
||||
|
||||
/// The native route relays bytes as they are. Python's synthetic `api_error` for a stream
|
||||
/// that never reaches `message_stop` lives in its SSE wrapper, above this route.
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_stream_that_ends_without_message_stop_is_relayed_as_is(call: MessagesCall) {
|
||||
const INCOMPLETE: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\n";
|
||||
let upstream =
|
||||
upstream([ResponseTemplate::new(200).set_body_raw(INCOMPLETE, "text/event-stream")]).await;
|
||||
let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX);
|
||||
|
||||
stream_through(&host).await.expect("streamed call succeeds");
|
||||
|
||||
let delivered: Vec<u8> = host
|
||||
.seen
|
||||
.into_inner()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.flat_map(|step| match step {
|
||||
Seen::Deliver(chunk) => chunk.to_vec(),
|
||||
Seen::Open(_) => Vec::new(),
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(delivered, INCOMPLETE.as_bytes());
|
||||
}
|
||||
|
||||
/// Serves one SSE chunk and then holds the connection open without ever finishing.
|
||||
async fn stalling_upstream() -> String {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base = format!("http://{}", listener.local_addr().unwrap());
|
||||
tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
let mut request = vec![0; 4096];
|
||||
let _ = socket.read(&mut request).await;
|
||||
socket
|
||||
.write_all(
|
||||
b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n\
|
||||
1f\r\nevent: message_start\ndata: {}\n\n\r\n",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
std::future::pending::<()>().await;
|
||||
});
|
||||
base
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) {
|
||||
let base = stalling_upstream().await;
|
||||
let host = RecordingStreamHost::new(
|
||||
MessagesCall {
|
||||
timeout: Some(Duration::from_millis(300)),
|
||||
..streaming(call, base)
|
||||
},
|
||||
usize::MAX,
|
||||
);
|
||||
|
||||
let error = tokio::time::timeout(Duration::from_secs(5), stream_through(&host))
|
||||
.await
|
||||
.expect("the stalled stream gives up within the timeout")
|
||||
.err()
|
||||
.expect("a stalled body fails the call");
|
||||
|
||||
assert!(matches!(error, Error::Transport(_)), "{error:?}");
|
||||
let seen = host.seen.into_inner().unwrap();
|
||||
assert!(
|
||||
matches!(seen.as_slice(), [Seen::Open(_), Seen::Deliver(chunk)] if chunk.as_ref() == b"event: message_start\ndata: {}\n\n"),
|
||||
"the chunk before the stall reached the caller, saw {} ops",
|
||||
seen.len()
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn streaming_is_refused_for_providers_that_cannot_stream(call: MessagesCall) {
|
||||
|
|
|
|||
|
|
@ -134,6 +134,13 @@ pub trait ProtocolHost: Send + Sync {
|
|||
response: <Self::Protocol as Protocol>::Response,
|
||||
) -> PyResult<Py<PyAny>>;
|
||||
|
||||
/// What the stream carries at hand-off, as the caller's stream receives it.
|
||||
fn head(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
head: <Self::Protocol as Protocol>::StreamHead,
|
||||
) -> PyResult<Py<PyAny>>;
|
||||
|
||||
/// One streamed chunk as the caller receives it.
|
||||
fn chunk(
|
||||
&mut self,
|
||||
|
|
|
|||
|
|
@ -134,10 +134,10 @@ where
|
|||
}
|
||||
match driver.resume(None)? {
|
||||
ExecutionStep::Return(value) => Ok(value),
|
||||
ExecutionStep::Open => py
|
||||
ExecutionStep::Open(head) => py
|
||||
.import("litellm.rust_bridge.lifecycle")?
|
||||
.getattr("SyncStream")?
|
||||
.call1((Py::new(py, Execution::suspended(driver))?,))
|
||||
.call1((Py::new(py, Execution::suspended(driver))?, head))
|
||||
.map(Bound::unbind),
|
||||
ExecutionStep::Await(_) | ExecutionStep::Yield(_) => {
|
||||
Err(PyRuntimeError::new_err("sync call suspended"))
|
||||
|
|
@ -312,7 +312,7 @@ where
|
|||
Ok(_) => return Err(missing_state()),
|
||||
Err(error) => Err(error),
|
||||
},
|
||||
HostOp::Open(_, reply) => return self.opened(py, reply).map(Next::Return),
|
||||
HostOp::Open(head, reply) => return self.opened(py, head, reply).map(Next::Return),
|
||||
HostOp::Deliver(chunk, reply) => {
|
||||
return self.delivered(py, chunk, reply).map(Next::Return);
|
||||
}
|
||||
|
|
@ -340,12 +340,21 @@ where
|
|||
}
|
||||
}
|
||||
|
||||
fn opened(&mut self, py: Python<'_>, reply: Reply<Demand>) -> PyResult<ExecutionStep> {
|
||||
fn opened(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
head: <ProtocolOf<H> as Protocol>::StreamHead,
|
||||
reply: Reply<Demand>,
|
||||
) -> PyResult<ExecutionStep> {
|
||||
self.stage = Stage::Streaming;
|
||||
let head = match self.host.head(py, head) {
|
||||
Ok(head) => head,
|
||||
Err(error) => return self.interrupt(py, error),
|
||||
};
|
||||
match self.adapter.opened(py) {
|
||||
Ok(()) => {
|
||||
self.pending = Some(Pending::Consumer(reply));
|
||||
Ok(ExecutionStep::Open)
|
||||
Ok(ExecutionStep::Open(head))
|
||||
}
|
||||
Err(error) => self.interrupt(py, error),
|
||||
}
|
||||
|
|
@ -699,6 +708,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
|
|||
.map(|answer| reply.send(answer))
|
||||
}
|
||||
|
||||
fn head(&mut self, _: Python<'_>, head: std::convert::Infallible) -> PyResult<Py<PyAny>> {
|
||||
match head {}
|
||||
}
|
||||
|
||||
fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult<Py<PyAny>> {
|
||||
match chunk {}
|
||||
}
|
||||
|
|
@ -945,6 +958,163 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
|
|||
});
|
||||
}
|
||||
|
||||
struct Streaming;
|
||||
|
||||
impl Protocol for Streaming {
|
||||
type Response = ();
|
||||
type Error = Error;
|
||||
type Projection = ();
|
||||
type Op = std::convert::Infallible;
|
||||
type Chunk = &'static str;
|
||||
type StreamHead = Vec<(&'static str, &'static str)>;
|
||||
}
|
||||
|
||||
struct StreamingHost;
|
||||
|
||||
impl ProtocolHost for StreamingHost {
|
||||
type Protocol = Streaming;
|
||||
type Failure = Classified;
|
||||
|
||||
fn project(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
_: &Bound<'_, PyDict>,
|
||||
) -> Result<(), InvokeError<Error>> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn invoke(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
op: std::convert::Infallible,
|
||||
) -> Result<(), InvokeError<Error>> {
|
||||
match op {}
|
||||
}
|
||||
|
||||
fn head(
|
||||
&mut self,
|
||||
py: Python<'_>,
|
||||
head: Vec<(&'static str, &'static str)>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let headers = PyDict::new(py);
|
||||
for (name, value) in head {
|
||||
headers.set_item(name, value)?;
|
||||
}
|
||||
let hidden = PyDict::new(py);
|
||||
hidden.set_item("additional_headers", headers)?;
|
||||
Ok(hidden.into_any().unbind())
|
||||
}
|
||||
|
||||
fn chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult<Py<PyAny>> {
|
||||
Ok(pyo3::types::PyString::new(py, chunk).into_any().unbind())
|
||||
}
|
||||
|
||||
fn complete(&mut self, py: Python<'_>, (): ()) -> PyResult<Py<PyAny>> {
|
||||
Ok(py.None())
|
||||
}
|
||||
|
||||
fn classify(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
|
||||
Ok(Classified(error.0))
|
||||
}
|
||||
|
||||
fn host_error(error: &PyErr) -> Error {
|
||||
Error(error.to_string())
|
||||
}
|
||||
|
||||
fn close(&mut self, _: Python<'_>) {}
|
||||
|
||||
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn streaming_machine() -> CallMachine<Streaming> {
|
||||
CallMachine::new(|host| {
|
||||
Box::pin(async move {
|
||||
host.project().await?;
|
||||
if host.open(vec![("request-id", "req_1")]).await? == Demand::Detached {
|
||||
return Ok(());
|
||||
}
|
||||
for chunk in ["first", "second"] {
|
||||
if host.deliver(chunk).await? == Demand::Detached {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// Drives a `Stream` (async) or `SyncStream` to completion from a sync test.
|
||||
fn read_all(py: Python<'_>, stream: &Bound<'_, PyAny>, asynchronous: bool) -> Vec<String> {
|
||||
if !asynchronous {
|
||||
return stream
|
||||
.try_iter()
|
||||
.unwrap()
|
||||
.map(|chunk| chunk.unwrap().extract().unwrap())
|
||||
.collect();
|
||||
}
|
||||
std::iter::from_fn(|| {
|
||||
let stop = stream
|
||||
.call_method0("__anext__")
|
||||
.unwrap()
|
||||
.call_method1("send", (py.None(),))
|
||||
.unwrap_err();
|
||||
if stop.is_instance_of::<pyo3::exceptions::PyStopAsyncIteration>(py) {
|
||||
return None;
|
||||
}
|
||||
assert!(stop.is_instance_of::<pyo3::exceptions::PyStopIteration>(py));
|
||||
Some(stop.value(py).getattr("value").unwrap().extract().unwrap())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_stream_carries_its_head_as_hidden_params_before_the_first_chunk() {
|
||||
let _guard = PYTHON_GLOBALS
|
||||
.lock()
|
||||
.unwrap_or_else(|error| error.into_inner());
|
||||
crate::initialize_python();
|
||||
Python::attach(|py| {
|
||||
install_lifecycle_module(py);
|
||||
for asynchronous in [false, true] {
|
||||
let log = Log::default();
|
||||
let adapter = SyntheticAdapter {
|
||||
log: Log(log.0.clone()),
|
||||
script: AdapterScript::Plain,
|
||||
};
|
||||
let handed = run_call(
|
||||
py,
|
||||
streaming_machine(),
|
||||
StreamingHost,
|
||||
Box::new(adapter),
|
||||
PyDict::new(py).unbind(),
|
||||
asynchronous,
|
||||
)
|
||||
.unwrap();
|
||||
let stream = if asynchronous {
|
||||
let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err();
|
||||
stop.value(py).getattr("value").unwrap()
|
||||
} else {
|
||||
handed.into_bound(py)
|
||||
};
|
||||
let hidden: std::collections::HashMap<
|
||||
String,
|
||||
std::collections::HashMap<String, String>,
|
||||
> = stream.getattr("_hidden_params").unwrap().extract().unwrap();
|
||||
assert_eq!(
|
||||
hidden["additional_headers"],
|
||||
std::collections::HashMap::from([(
|
||||
"request-id".to_string(),
|
||||
"req_1".to_string()
|
||||
)])
|
||||
);
|
||||
assert_eq!(log.entries(), ["started", "begin", "opened"]);
|
||||
assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn failing_machine() -> CallMachine<Synthetic> {
|
||||
CallMachine::new(|host| {
|
||||
Box::pin(async move {
|
||||
|
|
@ -1202,6 +1372,13 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
|
|||
) -> Result<(), InvokeError<Error>> {
|
||||
Err(missing_state().into())
|
||||
}
|
||||
fn head(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
head: std::convert::Infallible,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
match head {}
|
||||
}
|
||||
fn chunk(
|
||||
&mut self,
|
||||
_: Python<'_>,
|
||||
|
|
|
|||
|
|
@ -8,9 +8,9 @@ use pyo3::prelude::*;
|
|||
pub enum ExecutionStep {
|
||||
Return(Py<PyAny>),
|
||||
Await(Py<PyAny>),
|
||||
/// The call streams: the caller gets a stream over this execution, which stays
|
||||
/// suspended until the stream asks for a chunk.
|
||||
Open,
|
||||
/// The call streams: the caller gets a stream over this execution carrying this head,
|
||||
/// and the execution stays suspended until the stream asks for a chunk.
|
||||
Open(Py<PyAny>),
|
||||
Yield(Py<PyAny>),
|
||||
}
|
||||
|
||||
|
|
@ -75,7 +75,7 @@ impl Execution {
|
|||
let step = body.resume(result)?;
|
||||
let (tag, value, suspended) = match step {
|
||||
ExecutionStep::Await(value) => ("Await", value, true),
|
||||
ExecutionStep::Open => ("Open", py.None(), true),
|
||||
ExecutionStep::Open(head) => ("Open", head, true),
|
||||
ExecutionStep::Yield(value) => ("Yield", value, true),
|
||||
ExecutionStep::Return(value) => ("Complete", value, false),
|
||||
};
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ use std::convert::Infallible;
|
|||
use bytes::Bytes;
|
||||
use litellm_core::messages::{
|
||||
Error,
|
||||
route::{Messages, MessagesCall, MessagesOutput},
|
||||
route::{Messages, MessagesCall, MessagesOutput, MessagesStreamHead},
|
||||
types::MessagesShaping,
|
||||
};
|
||||
use litellm_host_python::{InvokeError, ProtocolHost, from_py, lookup, to_py};
|
||||
|
|
@ -238,6 +238,13 @@ impl ProtocolHost for MessagesPythonHost {
|
|||
}
|
||||
}
|
||||
|
||||
fn head(&mut self, py: Python<'_>, head: MessagesStreamHead) -> PyResult<Py<PyAny>> {
|
||||
py.import(ROUTE_HOST_MODULE)?
|
||||
.getattr("stream_hidden_params")?
|
||||
.call1((to_py(py, &head.headers)?,))
|
||||
.map(Bound::unbind)
|
||||
}
|
||||
|
||||
fn chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult<Py<PyAny>> {
|
||||
Ok(PyBytes::new(py, &chunk).into_any().unbind())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,14 +4,12 @@ use host::MessagesPythonHost;
|
|||
use litellm_callbacks_legacy_python::{
|
||||
LegacySurface, PassThroughStream, PublicCall, run_legacy_call,
|
||||
};
|
||||
use litellm_core::messages::route::{messages_machine, supports};
|
||||
use litellm_core::messages::route::messages_machine;
|
||||
use pyo3::{
|
||||
prelude::*,
|
||||
types::{PyDict, PyTuple},
|
||||
};
|
||||
|
||||
use crate::errors::RustBridgeDeclined;
|
||||
|
||||
const SURFACE: LegacySurface = LegacySurface {
|
||||
call_type: "anthropic_messages",
|
||||
input_description: "Messages",
|
||||
|
|
@ -28,17 +26,6 @@ fn run_messages(
|
|||
kwargs: Bound<'_, PyDict>,
|
||||
asynchronous: bool,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let model: String = request.getattr("model")?.extract()?;
|
||||
let provider: Option<String> = request.getattr("custom_llm_provider")?.extract()?;
|
||||
let stream = request
|
||||
.getattr("stream")?
|
||||
.extract::<Option<bool>>()?
|
||||
.unwrap_or(false);
|
||||
if !supports(&model, provider.as_deref(), stream) {
|
||||
return Err(RustBridgeDeclined::new_err(
|
||||
"the Rust Messages route does not serve this provider",
|
||||
));
|
||||
}
|
||||
let secrets = crate::secrets::source(py)?;
|
||||
run_legacy_call(
|
||||
py,
|
||||
|
|
|
|||
|
|
@ -117,6 +117,10 @@ impl ProtocolHost for OcrPythonHost {
|
|||
.map(Bound::unbind)
|
||||
}
|
||||
|
||||
fn head(&mut self, _: Python<'_>, head: std::convert::Infallible) -> PyResult<Py<PyAny>> {
|
||||
match head {}
|
||||
}
|
||||
|
||||
fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult<Py<PyAny>> {
|
||||
match chunk {}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Itera
|
|||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
|
||||
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages import handler as main
|
||||
from litellm.rust_bridge.catalog import Delivery, Route, RouteContext
|
||||
from litellm.rust_bridge.dispatch import PublicDispatch, call_hook
|
||||
|
|
@ -71,10 +73,17 @@ def _public_request(
|
|||
)
|
||||
|
||||
|
||||
def _resolved_provider(request: LiteLLMMessagesRequest) -> str | None:
|
||||
try:
|
||||
return get_llm_provider(request.model, request.custom_llm_provider)[1]
|
||||
except BadRequestError:
|
||||
return request.custom_llm_provider
|
||||
|
||||
|
||||
def _context(request: LiteLLMMessagesRequest) -> RouteContext:
|
||||
return RouteContext(
|
||||
Route.MESSAGES,
|
||||
provider=request.custom_llm_provider,
|
||||
provider=_resolved_provider(request),
|
||||
model=request.model,
|
||||
delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -109,6 +109,7 @@ RULES: Final[Rules] = (
|
|||
RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY),
|
||||
RouteRule(Route.EMBEDDINGS, Rollout.PYTHON_ONLY),
|
||||
RouteRule(Route.OCR, Rollout.RUST_REQUIRED),
|
||||
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, providers=frozenset({"anthropic"})),
|
||||
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY),
|
||||
RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY),
|
||||
RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY),
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Awaitable, Iterator
|
||||
from collections.abc import AsyncIterator, Awaitable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Protocol
|
||||
|
||||
|
|
@ -17,7 +17,7 @@ class Complete:
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Open:
|
||||
value: None
|
||||
value: Mapping[str, object] | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -68,7 +68,7 @@ async def drive(execution: Execution) -> object:
|
|||
step: Final = await _settle(execution, execution.start())
|
||||
if isinstance(step, Open):
|
||||
handed_off = True
|
||||
return Stream(execution)
|
||||
return Stream(execution, step.value)
|
||||
return step.value
|
||||
finally:
|
||||
if not handed_off:
|
||||
|
|
@ -78,10 +78,10 @@ async def drive(execution: Execution) -> object:
|
|||
class Stream(AsyncIterator[object]):
|
||||
"""A streamed native call: each read resumes the execution until its next chunk."""
|
||||
|
||||
def __init__(self, execution: Execution) -> None:
|
||||
def __init__(self, execution: Execution, hidden_params: Mapping[str, object] | None = None) -> None:
|
||||
self._execution: Final = execution
|
||||
self._done = False
|
||||
self._hidden_params: dict[str, object] = {} # mutable-ok: header writers mutate _hidden_params in place
|
||||
self._hidden_params: dict[str, object] = dict(hidden_params or {}) # mutable-ok: header writers mutate it
|
||||
|
||||
def __aiter__(self) -> Stream:
|
||||
return self
|
||||
|
|
@ -115,10 +115,10 @@ class Stream(AsyncIterator[object]):
|
|||
class SyncStream(Iterator[object]):
|
||||
"""The sync form of `Stream`; its execution never suspends on an awaitable."""
|
||||
|
||||
def __init__(self, execution: Execution) -> None:
|
||||
def __init__(self, execution: Execution, hidden_params: Mapping[str, object] | None = None) -> None:
|
||||
self._execution: Final = execution
|
||||
self._done = False
|
||||
self._hidden_params: dict[str, object] = {} # mutable-ok: header writers mutate _hidden_params in place
|
||||
self._hidden_params: dict[str, object] = dict(hidden_params or {}) # mutable-ok: header writers mutate it
|
||||
|
||||
def __iter__(self) -> SyncStream:
|
||||
return self
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from collections.abc import Mapping, Sequence
|
|||
from dataclasses import asdict, dataclass
|
||||
from typing import Final, cast # noqa: TID251 # narrows the normalized native payload to the public TypedDict
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
|
|
@ -53,6 +54,14 @@ def response(value: Mapping[str, object]) -> AnthropicMessagesResponse:
|
|||
)
|
||||
|
||||
|
||||
def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, object]:
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
anthropic_messages_stream_hidden_params,
|
||||
)
|
||||
|
||||
return anthropic_messages_stream_hidden_params(httpx.Headers(list(headers)))
|
||||
|
||||
|
||||
def arguments(request: LiteLLMMessagesRequest) -> Mapping[str, object]:
|
||||
return request.kwargs
|
||||
|
||||
|
|
|
|||
|
|
@ -110,3 +110,15 @@ def test_native_request_rejections_map_to_the_public_400() -> None:
|
|||
assert "does not support top_k=5" in mapped.message
|
||||
assert mapped.model == "claude-sonnet-5"
|
||||
assert not isinstance(route_host.map_failure(ValueError("plain"), request, "anthropic"), litellm.BadRequestError)
|
||||
|
||||
|
||||
def test_stream_hidden_params_projects_upstream_headers_the_way_the_python_handler_does() -> None:
|
||||
hidden: Final = route_host.stream_hidden_params(
|
||||
(("request-id", "req_upstream_123"), ("x-ratelimit-remaining-requests", "41"))
|
||||
)
|
||||
|
||||
additional: Final = hidden["additional_headers"]
|
||||
assert isinstance(additional, dict)
|
||||
assert additional["llm_provider-request-id"] == "req_upstream_123"
|
||||
assert additional["x-ratelimit-remaining-requests"] == "41"
|
||||
assert "request-id" not in additional
|
||||
|
|
|
|||
|
|
@ -127,7 +127,7 @@ async def test_native_messages_stream_relays_provider_events_and_logs_success_on
|
|||
**arguments(messages_server, stream=True, callbacks=[recorder])
|
||||
)
|
||||
assert isinstance(stream, AsyncIterator)
|
||||
assert get_hidden_params_dict(stream) == {"additional_headers": {"x-litellm-rust": "true"}}
|
||||
assert get_hidden_params_dict(stream)["additional_headers"]["x-litellm-rust"] == "true"
|
||||
first: Final = await anext(stream)
|
||||
await drain_logging()
|
||||
assert "async_log_success_event" not in recorder.names
|
||||
|
|
@ -171,7 +171,7 @@ def test_native_sync_messages_stream_relays_provider_events_and_logs_success_onc
|
|||
|
||||
stream: Final = litellm.anthropic.messages.create(**arguments(messages_server, stream=True, callbacks=[recorder]))
|
||||
assert isinstance(stream, Iterator)
|
||||
assert get_hidden_params_dict(stream) == {"additional_headers": {"x-litellm-rust": "true"}}
|
||||
assert get_hidden_params_dict(stream)["additional_headers"]["x-litellm-rust"] == "true"
|
||||
|
||||
assert b"".join(stream) == sse_payload()
|
||||
assert_served_natively(messages_server)
|
||||
|
|
@ -186,3 +186,56 @@ def test_native_sync_messages_returns_the_provider_message(messages_server: Reco
|
|||
assert_served_natively(messages_server)
|
||||
assert response["content"] == MESSAGES_RESPONSE["content"]
|
||||
assert len(recorder.wait_for("log_success_event")) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_messages_pre_call_sees_the_shaped_optional_params(
|
||||
messages_server: RecordingServer,
|
||||
) -> None:
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
await litellm.anthropic.messages.acreate(
|
||||
**arguments(messages_server, callbacks=[recorder], temperature=0.2, top_k=3, drop_params=True)
|
||||
)
|
||||
|
||||
sent: Final = messages_server.requests[0].body
|
||||
assert not {"temperature", "top_k"} & sent.keys()
|
||||
pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs
|
||||
assert isinstance(pre_call, dict)
|
||||
optional_params: Final = pre_call["optional_params"]
|
||||
assert isinstance(optional_params, dict)
|
||||
assert not {"model", "messages", "temperature", "top_k"} & optional_params.keys()
|
||||
assert optional_params["max_tokens"] == sent["max_tokens"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_messages_failing_pre_call_logger_does_not_fail_the_call(messages_server: RecordingServer) -> None:
|
||||
class Broken(CustomLogger):
|
||||
def log_pre_api_call(self, model, messages, kwargs):
|
||||
raise RuntimeError("logger exploded")
|
||||
|
||||
response: Final = await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Broken()]))
|
||||
|
||||
assert_served_natively(messages_server)
|
||||
assert response["content"] == MESSAGES_RESPONSE["content"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_messages_stream_success_log_carries_usage_rebuilt_from_the_relayed_events(
|
||||
messages_server: RecordingServer,
|
||||
) -> None:
|
||||
messages_server.enqueue(STREAM)
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
stream: Final = await litellm.anthropic.messages.acreate(
|
||||
**arguments(messages_server, stream=True, callbacks=[recorder])
|
||||
)
|
||||
assert isinstance(stream, AsyncIterator)
|
||||
async for _ in stream:
|
||||
pass
|
||||
|
||||
success: Final = await recorder.wait_for_async("async_log_success_event")
|
||||
usage: Final = success[0].response.usage
|
||||
assert usage.completion_tokens == MESSAGES_EVENTS[4][1]["usage"]["output_tokens"]
|
||||
assert usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"]
|
||||
assert success[0].response.choices[0].message.content == "Hello from native Messages"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue