mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
test(rust): add Messages live provider harness (#45823)
* test(rust): add Messages live provider harness * test(rust): support automatic Vertex tool selection
This commit is contained in:
parent
c2bc5aecff
commit
8db9037c90
9 changed files with 440 additions and 0 deletions
3
litellm-rust/Cargo.lock
generated
3
litellm-rust/Cargo.lock
generated
|
|
@ -4024,6 +4024,7 @@ dependencies = [
|
|||
"litellm-cache-memory",
|
||||
"litellm-cache-response",
|
||||
"litellm-core-utils",
|
||||
"litellm-framer",
|
||||
"litellm-host",
|
||||
"litellm-host-native",
|
||||
"litellm-http",
|
||||
|
|
@ -4107,9 +4108,11 @@ name = "litellm-inference-testing"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"litellm-auth",
|
||||
"litellm-http",
|
||||
"litellm-inference",
|
||||
"litellm-secrets",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ serde_json = { workspace = true, features = ["preserve_order"] }
|
|||
tracing.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-framer = { workspace = true, features = ["sse"] }
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-host-native.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
|
|
|
|||
5
litellm-rust/crates/inference-messages/tests/live.rs
Normal file
5
litellm-rust/crates/inference-messages/tests/live.rs
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
#[path = "live/support.rs"]
|
||||
mod support;
|
||||
|
||||
#[path = "live/vertex_ai.rs"]
|
||||
mod vertex_ai;
|
||||
265
litellm-rust/crates/inference-messages/tests/live/support.rs
Normal file
265
litellm-rust/crates/inference-messages/tests/live/support.rs
Normal file
|
|
@ -0,0 +1,265 @@
|
|||
use std::{sync::OnceLock, time::Duration};
|
||||
|
||||
use futures_util::{FutureExt, TryStreamExt};
|
||||
use litellm_host::interceptors::{
|
||||
ExecutionFacts, Interceptors, RawResponse, RequestContext, ResultSource, WireRequest,
|
||||
};
|
||||
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesRoute};
|
||||
use litellm_inference_testing::live::{LiveResources, model, within_deadline};
|
||||
use litellm_llms_types::formats::messages::MessagesResponse;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
pub(super) struct LiveCall {
|
||||
provider: &'static str,
|
||||
facts: OnceLock<ExecutionFacts>,
|
||||
}
|
||||
|
||||
impl LiveCall {
|
||||
pub(super) fn new(provider: &'static str) -> Self {
|
||||
Self {
|
||||
provider,
|
||||
facts: OnceLock::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn facts(&self) -> &ExecutionFacts {
|
||||
self.facts
|
||||
.get()
|
||||
.expect("the route must deliver execution facts")
|
||||
}
|
||||
|
||||
pub(super) fn assert_provider_result(&self) {
|
||||
let facts = self.facts();
|
||||
assert_eq!(facts.provider.provider, self.provider);
|
||||
assert_eq!(facts.source, ResultSource::Provider);
|
||||
println!(
|
||||
"{}",
|
||||
json!({"route": "messages", "provider": self.provider, "source": "provider"})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
impl Interceptors<Error> for LiveCall {
|
||||
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> {
|
||||
assert!(
|
||||
self.facts.set(facts).is_ok(),
|
||||
"result must be accepted exactly once"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn before_provider_request(
|
||||
&self,
|
||||
wire: WireRequest,
|
||||
context: RequestContext,
|
||||
) -> Result<WireRequest, Error> {
|
||||
assert_eq!(context.custom_llm_provider, self.provider);
|
||||
println!(
|
||||
"{}",
|
||||
json!({"route": "messages", "provider": self.provider, "model": context.model})
|
||||
);
|
||||
Ok(wire)
|
||||
}
|
||||
|
||||
async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::fixture]
|
||||
pub(super) fn route() -> MessagesRoute {
|
||||
let resources = LiveResources::default();
|
||||
MessagesRoute::new(resources.http, resources.auth, resources.secrets)
|
||||
}
|
||||
|
||||
pub(super) async fn request(
|
||||
route: &MessagesRoute,
|
||||
call: MessagesCall,
|
||||
host: &LiveCall,
|
||||
) -> MessagesCallResponse {
|
||||
route
|
||||
.execute(call, host, None)
|
||||
.boxed()
|
||||
.await
|
||||
.expect("live request failed")
|
||||
}
|
||||
|
||||
pub(super) fn call(provider: &str, payload: Value) -> MessagesCall {
|
||||
let Value::Object(fields) = payload else {
|
||||
panic!("live payload must be an object")
|
||||
};
|
||||
let body = serde_json::from_value(Value::Object(
|
||||
[("max_tokens".into(), json!(128))]
|
||||
.into_iter()
|
||||
.chain(fields)
|
||||
.chain([("model".into(), json!(model("messages", provider)))])
|
||||
.collect(),
|
||||
))
|
||||
.expect("live payload must be a valid Messages request");
|
||||
MessagesCall {
|
||||
body,
|
||||
api_key: None,
|
||||
api_base: None,
|
||||
custom_llm_provider: Some(provider.into()),
|
||||
litellm_params: Default::default(),
|
||||
extra_headers: None,
|
||||
provider_specific_header: None,
|
||||
timeout: Some(Duration::from_secs(60)),
|
||||
shaping: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn complete(response: MessagesCallResponse) -> Box<MessagesResponse> {
|
||||
match response {
|
||||
MessagesCallResponse::Complete(message) => message,
|
||||
MessagesCallResponse::Stream { .. } => panic!("a complete request returned a stream"),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn assert_text(message: &MessagesResponse) {
|
||||
assert!(!message.id.is_empty());
|
||||
assert_eq!(message.role, "assistant");
|
||||
assert!(message.stop_reason.is_some());
|
||||
assert!(
|
||||
message
|
||||
.content
|
||||
.iter()
|
||||
.any(|block| block["text"].as_str().is_some_and(|text| !text.is_empty()))
|
||||
);
|
||||
assert!(
|
||||
message
|
||||
.usage
|
||||
.as_ref()
|
||||
.and_then(|usage| usage["output_tokens"].as_u64())
|
||||
.is_some_and(|count| count > 0)
|
||||
);
|
||||
}
|
||||
|
||||
pub(super) async fn stream_events(response: MessagesCallResponse) -> Vec<Value> {
|
||||
let MessagesCallResponse::Stream { head, chunks } = response else {
|
||||
panic!("stream request returned a complete response")
|
||||
};
|
||||
let content_types: Vec<_> = head
|
||||
.headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("content-type"))
|
||||
.map(|(_, value)| value.split(';').next().unwrap().trim())
|
||||
.collect();
|
||||
assert_eq!(content_types, ["text/event-stream"]);
|
||||
let frames = litellm_framer::frames(chunks, litellm_framer::sse::SseCodec::default())
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.expect("live SSE framing failed");
|
||||
let data_frames = match frames.last() {
|
||||
Some(frame) if frame.data == "[DONE]" => &frames[..frames.len() - 1],
|
||||
_ => frames.as_slice(),
|
||||
};
|
||||
let events: Vec<Value> = data_frames
|
||||
.iter()
|
||||
.map(|frame| serde_json::from_str(&frame.data).expect("SSE data must be JSON"))
|
||||
.collect();
|
||||
assert_eq!(
|
||||
events.first().map(|event| event["type"].as_str()),
|
||||
Some(Some("message_start"))
|
||||
);
|
||||
assert_eq!(
|
||||
events.last().map(|event| event["type"].as_str()),
|
||||
Some(Some("message_stop"))
|
||||
);
|
||||
assert!(!events.iter().any(|event| event["type"] == "error"));
|
||||
assert!(events.iter().any(|event| {
|
||||
event["usage"]["output_tokens"]
|
||||
.as_u64()
|
||||
.is_some_and(|count| count > 0)
|
||||
}));
|
||||
println!("{}", json!({"events": events}));
|
||||
events
|
||||
}
|
||||
|
||||
fn prompt() -> Value {
|
||||
json!({"role": "user", "content": "Use the echo tool with value live-check, then tell me the result."})
|
||||
}
|
||||
|
||||
fn tools() -> Value {
|
||||
json!([{"name": "echo", "description": "Return a value", "input_schema": {
|
||||
"type": "object", "properties": {"value": {"type": "string", "enum": ["live-check"]}}, "required": ["value"]
|
||||
}}])
|
||||
}
|
||||
|
||||
fn fragments(events: &[Value], index: u64, field: &str) -> String {
|
||||
events
|
||||
.iter()
|
||||
.filter(|event| event["index"].as_u64() == Some(index))
|
||||
.filter_map(|event| event["delta"][field].as_str())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn streamed_block(start: &Value, events: &[Value]) -> Value {
|
||||
let index = start["index"]
|
||||
.as_u64()
|
||||
.expect("content block must have an index");
|
||||
let block = &start["content_block"];
|
||||
match block["type"].as_str() {
|
||||
Some("tool_use") => {
|
||||
let partial = fragments(events, index, "partial_json");
|
||||
let input = match partial.is_empty() {
|
||||
true => block["input"].clone(),
|
||||
false => {
|
||||
serde_json::from_str(&partial).expect("streamed tool arguments must be JSON")
|
||||
}
|
||||
};
|
||||
json!({"type": "tool_use", "id": block["id"], "name": block["name"], "input": input})
|
||||
}
|
||||
Some("text") => {
|
||||
json!({"type": "text", "text": format!("{}{}", block["text"].as_str().unwrap_or_default(), fragments(events, index, "text"))})
|
||||
}
|
||||
Some("thinking") => {
|
||||
json!({"type": "thinking", "thinking": format!("{}{}", block["thinking"].as_str().unwrap_or_default(), fragments(events, index, "thinking")),
|
||||
"signature": format!("{}{}", block["signature"].as_str().unwrap_or_default(), fragments(events, index, "signature"))})
|
||||
}
|
||||
_ => block.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn tool_round_trip(
|
||||
route: MessagesRoute,
|
||||
provider: &'static str,
|
||||
stream: bool,
|
||||
tool_params: Value,
|
||||
) {
|
||||
within_deadline(async {
|
||||
let first_host = LiveCall::new(provider);
|
||||
let payload = json!({"messages": [prompt()], "tools": tools(), "stream": stream});
|
||||
let fields = payload.as_object().expect("tool payload must be an object").clone();
|
||||
let params = tool_params.as_object().expect("tool parameters must be an object").clone();
|
||||
let first = request(&route, call(provider, Value::Object(fields.into_iter().chain(params).collect())), &first_host).await;
|
||||
let content = match stream {
|
||||
false => {
|
||||
let message = complete(first);
|
||||
assert_eq!(message.stop_reason.as_deref(), Some("tool_use"));
|
||||
message.content
|
||||
}
|
||||
true => {
|
||||
let events = stream_events(first).await;
|
||||
assert!(events.iter().any(|event| event["delta"]["stop_reason"] == "tool_use"));
|
||||
events.iter().filter(|event| event["type"] == "content_block_start")
|
||||
.map(|event| streamed_block(event, &events)).collect()
|
||||
}
|
||||
};
|
||||
let tool = content.iter().find(|block| block["type"] == "tool_use").expect("requested tool must be returned");
|
||||
assert_eq!(tool["name"], "echo");
|
||||
assert_eq!(tool["input"]["value"], "live-check");
|
||||
let tool_id = tool["id"].as_str().filter(|id| !id.is_empty()).expect("tool must have an id");
|
||||
first_host.assert_provider_result();
|
||||
let second_host = LiveCall::new(provider);
|
||||
let second = complete(request(&route, call(provider, json!({
|
||||
"tools": tools(), "messages": [prompt(),
|
||||
{"role": "assistant", "content": content},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": tool_id, "content": "live-check"}]}]
|
||||
})), &second_host).await);
|
||||
assert_text(&second);
|
||||
second_host.assert_provider_result();
|
||||
println!("{}", serde_json::to_string(&second).unwrap());
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
|
@ -0,0 +1,68 @@
|
|||
use litellm_inference_messages::MessagesRoute;
|
||||
use litellm_inference_testing::live::within_deadline;
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::support::{LiveCall, assert_text, call, complete, request, route, stream_events};
|
||||
|
||||
const PROVIDER: &str = "vertex_ai";
|
||||
|
||||
#[rstest]
|
||||
#[case::basic(json!({"messages": [{"role": "user", "content": "Say hello."}]}))]
|
||||
#[case::system_text(json!({
|
||||
"system": "You are a helpful assistant.",
|
||||
"messages": [{"role": "user", "content": "Say hello."}]
|
||||
}))]
|
||||
#[ignore = "calls a real provider, requires credentials and a live model"]
|
||||
#[tokio::test]
|
||||
async fn completion(route: MessagesRoute, #[case] payload: Value) {
|
||||
within_deadline(async {
|
||||
let host = LiveCall::new(PROVIDER);
|
||||
let message = complete(request(&route, call(PROVIDER, payload), &host).await);
|
||||
assert_text(&message);
|
||||
host.assert_provider_result();
|
||||
println!("{}", serde_json::to_string(&message).unwrap());
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[ignore = "calls a real provider, requires credentials and a live model"]
|
||||
#[tokio::test]
|
||||
async fn streaming(route: MessagesRoute) {
|
||||
within_deadline(async {
|
||||
let host = LiveCall::new(PROVIDER);
|
||||
let response = request(
|
||||
&route,
|
||||
call(
|
||||
PROVIDER,
|
||||
json!({"stream": true, "messages": [{"role": "user", "content": "Say hello."}]}),
|
||||
),
|
||||
&host,
|
||||
)
|
||||
.await;
|
||||
let events = stream_events(response).await;
|
||||
assert!(events.iter().any(|event| {
|
||||
event["delta"]["text"]
|
||||
.as_str()
|
||||
.is_some_and(|text| !text.is_empty())
|
||||
}));
|
||||
host.assert_provider_result();
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::complete(false)]
|
||||
#[case::streaming(true)]
|
||||
#[ignore = "calls a real provider, requires credentials and a live model"]
|
||||
#[tokio::test]
|
||||
async fn tool_round_trip(route: MessagesRoute, #[case] stream: bool) {
|
||||
super::support::tool_round_trip(
|
||||
route,
|
||||
PROVIDER,
|
||||
stream,
|
||||
json!({"tool_choice": {"type": "auto"}}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
|
@ -7,7 +7,9 @@ repository.workspace = true
|
|||
publish = false
|
||||
|
||||
[dependencies]
|
||||
litellm-auth.workspace = true
|
||||
futures-util.workspace = true
|
||||
litellm-http.workspace = true
|
||||
litellm-inference.workspace = true
|
||||
litellm-secrets.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
41
litellm-rust/crates/inference-testing/README.md
Normal file
41
litellm-rust/crates/inference-testing/README.md
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
# Live provider route tests
|
||||
|
||||
Each route crate owns `tests/live.rs`. Use named `rstest` provider cases for the behaviors that route supports. Shared setup lives in `litellm_inference_testing::live`, with production HTTP settings, authentication services, environment secrets and a deadline
|
||||
|
||||
These tests execute the Rust route against a real provider. They do not prove Python parity, proxy authorization, billing persistence or agent compatibility. Those belong in the corresponding boundary and gateway suites
|
||||
|
||||
All live cases use `#[ignore = "calls a real provider, requires credentials and a live model"]`. Normal Cargo runs report them as ignored. An explicitly selected case fails on missing credentials, unavailable models, provider errors or timeouts. Never return early and count missing coverage as a pass
|
||||
|
||||
Set `LITELLM_LIVE_<ROUTE>_<PROVIDER>_MODEL` to the latest model available to the test account. The model has no compiled default because vendor availability changes independently of LiteLLM. Credentials use the provider's existing configuration rather than a second test-only key convention. Export credentials in the shell before running tests. The harness does not load `.env` or modify process environment while tests run
|
||||
|
||||
## Messages
|
||||
|
||||
Completion checks text and usage. Streaming consumes SSE to completion and checks text, usage and the terminal event. Tool round trips request a tool call and send its returned ID back with a result. Each case checks provider execution facts and rejects cache results
|
||||
|
||||
Messages setup and assertions live in `inference-messages/tests/live/support.rs`. Provider modules under `tests/live/` select their payloads and tool policy. The shared helpers consume streams completely and reconstruct text, thinking signatures and tool arguments for the follow-up request
|
||||
|
||||
The initial provider case is `vertex_ai`. A registered ignored case is available coverage, not evidence that it passed against that provider
|
||||
|
||||
Run one provider's cases explicitly
|
||||
|
||||
```bash
|
||||
LITELLM_LIVE_MESSAGES_VERTEX_AI_MODEL=<claude-model-on-vertex> \
|
||||
cargo test --manifest-path litellm-rust/Cargo.toml \
|
||||
-p litellm-inference-messages --test live vertex_ai \
|
||||
-- --ignored --nocapture --test-threads=1
|
||||
```
|
||||
|
||||
For Vertex AI, set `LITELLM_LIVE_MESSAGES_VERTEX_AI_MODEL` to a Claude model published in the project's location. Credentials, project and location resolve through the production Rust GCP authentication, from `GOOGLE_APPLICATION_CREDENTIALS` or `VERTEXAI_CREDENTIALS`, with `VERTEXAI_PROJECT` and `VERTEXAI_LOCATION`
|
||||
|
||||
Vertex AI cases live in `inference-messages/tests/live/vertex_ai.rs`. They cover plain and system messages, text streaming, and complete and streaming tool round trips. Tool cases use automatic selection because current Claude models reject forced `tool_choice`. Each case still requires a tool call, checks its arguments, and sends the tool result back
|
||||
|
||||
Tool cases require a model that supports forced `tool_choice`, except where a provider's section says otherwise. A provider rejection fails the selected case instead of skipping it
|
||||
|
||||
Select one behavior by its test name, for example `completion::case_1_basic`. To list every available cell without contacting any provider
|
||||
|
||||
```bash
|
||||
cargo test --manifest-path litellm-rust/Cargo.toml \
|
||||
-p litellm-inference-messages --test live -- --list
|
||||
```
|
||||
|
||||
For another route, reuse `LiveResources`, `model` and `within_deadline`. Keep payload construction and response assertions in that route's test file. Add only supported provider behaviors, and document any gap instead of substituting mocks or skipping failures. Run cases serially to keep spending and rate limits predictable
|
||||
|
|
@ -9,6 +9,8 @@ use litellm_secrets::{SecretValue, source::SecretSource};
|
|||
|
||||
use litellm_inference::resources::CoreResources;
|
||||
|
||||
pub mod live;
|
||||
|
||||
pub fn http_pool() -> HttpClientPool {
|
||||
HttpClientPool::new(Arc::new(PublicDnsResolver))
|
||||
}
|
||||
|
|
|
|||
53
litellm-rust/crates/inference-testing/src/live.rs
Normal file
53
litellm-rust/crates/inference-testing/src/live.rs
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
use std::{future::Future, sync::Arc, time::Duration};
|
||||
|
||||
use litellm_http::{ClientVariant, HttpSettings, HttpSettingsLayer, Resolution};
|
||||
use litellm_inference::resources::CoreResources;
|
||||
use litellm_secrets::source::{EnvironmentSecrets, SecretSource};
|
||||
|
||||
pub const DEADLINE: Duration = Duration::from_secs(120);
|
||||
|
||||
pub struct LiveResources {
|
||||
pub http: litellm_http::Client,
|
||||
pub auth: Arc<litellm_auth::AuthServices>,
|
||||
pub secrets: Arc<dyn SecretSource>,
|
||||
}
|
||||
|
||||
impl Default for LiveResources {
|
||||
fn default() -> Self {
|
||||
let resources = CoreResources::new(Arc::new(super::http_pool()));
|
||||
let settings =
|
||||
HttpSettings::from_layers([HttpSettingsLayer::from_environment(&|key: &str| {
|
||||
std::env::var(key).ok()
|
||||
})]);
|
||||
let resolution = Resolution::from(&settings);
|
||||
let http = resources
|
||||
.pool
|
||||
.client(&resolution.config, ClientVariant::Provider)
|
||||
.expect("live provider HTTP settings must be valid");
|
||||
let secrets_http = resources
|
||||
.pool
|
||||
.client(&resolution.config, ClientVariant::NoRedirect)
|
||||
.expect("live secret HTTP settings must be valid");
|
||||
Self {
|
||||
http,
|
||||
auth: resources.auth,
|
||||
secrets: Arc::new(EnvironmentSecrets::python_compatible(secrets_http)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn model(route: &str, provider: &str) -> String {
|
||||
let name = format!("LITELLM_LIVE_{route}_{provider}_MODEL").to_ascii_uppercase();
|
||||
std::env::var(&name)
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.unwrap_or_else(|| {
|
||||
panic!("set {name} to an available current model before running this live case")
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn within_deadline<T>(operation: impl Future<Output = T>) -> T {
|
||||
tokio::time::timeout(DEADLINE, operation)
|
||||
.await
|
||||
.expect("live case exceeded its 120 second deadline")
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue