diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index bc133ec2843..e645db0d4b5 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3997,6 +3997,7 @@ dependencies = [ name = "litellm-inference-chat" version = "0.1.0" dependencies = [ + "bytes", "futures-util", "litellm-auth", "litellm-cache-memory", @@ -4011,6 +4012,7 @@ dependencies = [ "litellm-llms-types", "litellm-secrets", "litellm-tracing", + "reqwest 0.12.28", "rstest", "serde_json", "tokio", diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 742a373b286..8b5e11c954f 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -3,13 +3,20 @@ use std::sync::Arc; use axum::{ Json, + body::Bytes, extract::{Path, State}, + http::header::CONTENT_TYPE, response::{IntoResponse, Response}, }; -use litellm_inference_chat::types::ChatCompletionsCall; +use litellm_host_http::Sse; +use litellm_inference_chat::{ + ChatCompletionsRoute, + openai::{OpenAiChatCall, OpenAiChatCompletions}, + types::ChatCompletionsCall, +}; use serde_json::{Map, Value}; -use crate::{Error, Gateway, JsonObject, request}; +use crate::{Deployment, Error, Gateway, JsonObject, request}; pub(crate) async fn create( State(gateway): State>, @@ -43,6 +50,12 @@ async fn handle( let deployment = request::resolve_deployment(gateway, &body)?; request::authorize_model(identity, deployment, &body).await?; let (body, cache_options) = crate::caching::prepare(identity, body)?; + if ChatCompletionsRoute::serves_openai( + &deployment.model, + deployment.custom_llm_provider.as_deref(), + ) { + return openai(gateway, deployment, body).await; + } let route = gateway.chat_completions.clone(); let route = match &gateway.cache { Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new( @@ -79,3 +92,27 @@ async fn handle( .await?; Ok(headers.apply(response)) } + +async fn openai( + gateway: &Gateway, + deployment: &Deployment, + body: Map, +) -> Result { + let machine = gateway.chat_completions.clone().openai_machine( + OpenAiChatCall { + model: deployment.model.clone(), + body, + api_key: deployment.api_key.clone(), + api_base: deployment.api_base.clone(), + custom_llm_provider: deployment.custom_llm_provider.clone(), + extra_headers: None, + timeout: deployment.timeout, + }, + None, + ); + let encoder = Sse::::new( + |body: Bytes| ([(CONTENT_TYPE, "application/json")], body), + |error| Bytes::from(Error::from(error).openai_sse_frame()), + ); + Ok(litellm_host_http::serve(machine, (), (), encoder, None).await?) +} diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 39656b933eb..3588c03e990 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -76,21 +76,30 @@ impl Error { } pub fn openai_response(self) -> Response { + (self.status(), Json(self.openai_body())).into_response() + } + + /// An OpenAI error event, for a Chat Completions stream that fails after its headers + /// went out. + pub fn openai_sse_frame(&self) -> String { + format!("data: {}\n\n", self.openai_body()) + } + + fn openai_body(&self) -> Value { let status = self.status(); - let message = match &self { + let message = match self { Self::UnknownModel(model) => format!("Invalid model name passed in model={model}"), + Self::Route(RouteError::Transport(TransportError::Http { body, .. })) => { + openai_error_message(body).unwrap_or_else(|| self.to_string()) + } _ => self.to_string(), }; - ( - status, - Json(json!({"error": { - "message": message, - "type": error_type(status), - "param": null, - "code": status.as_u16(), - }})), - ) - .into_response() + json!({"error": { + "message": message, + "type": error_type(status), + "param": null, + "code": status.as_u16(), + }}) } /// The Anthropic error envelope Python's `AnthropicExceptionMapping` builds: an upstream @@ -148,6 +157,14 @@ fn is_anthropic_error(object: &Map) -> bool { .is_some_and(|error| error.contains_key("type") && error.contains_key("message")) } +fn openai_error_message(body: &str) -> Option { + let body: Value = serde_json::from_str(body).ok()?; + body.get("error")? + .get("message")? + .as_str() + .map(str::to_owned) +} + fn provider_message(object: &Map) -> Option<&str> { if let Some(detail) = object.get("detail").and_then(Value::as_object) { return detail.get("message").and_then(Value::as_str); @@ -223,6 +240,24 @@ mod tests { assert_eq!(error.body(request_id), expected); } + #[rstest] + #[case::openai_body( + upstream(429, r#"{"error":{"message":"slow down","type":"requests","code":"rate_limit"}}"#), + json!({"error": {"message": "slow down", "type": "rate_limit_error", "param": null, "code": 429}}), + )] + #[case::plain_text( + upstream(502, "bad gateway"), + json!({"error": { + "message": "upstream request failed with status 502: bad gateway", + "type": "api_error", "param": null, "code": 502, + }}), + )] + fn openai_errors_carry_the_provider_message(#[case] error: Error, #[case] expected: Value) { + let frame = error.openai_sse_frame(); + assert_eq!(error.openai_body(), expected); + assert_eq!(frame, format!("data: {expected}\n\n")); + } + #[rstest] #[case::upstream_status(upstream(429, ""), StatusCode::TOO_MANY_REQUESTS)] #[case::rejected_request(Error::Route(RouteError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)] diff --git a/litellm-rust/crates/gateway-inference/tests/openai_chat.rs b/litellm-rust/crates/gateway-inference/tests/openai_chat.rs new file mode 100644 index 00000000000..4012513fbc3 --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/openai_chat.rs @@ -0,0 +1,190 @@ +mod support; + +use std::time::Duration; + +use axum::body::to_bytes; +use futures_util::StreamExt; +use rstest::rstest; +use serde_json::{Value, json}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, header, method, path}, +}; + +const COMPLETION: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-6-astra","choices":[{"index":0,"message":{"role":"assistant","content":"hi","refusal":null},"logprobs":null,"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4},"system_fingerprint":"fp"}"#; +const STREAM: &str = "data: {\"id\":\"c\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\ndata: {\"id\":\"c\",\"object\":\"chat.completion.chunk\",\"choices\":[],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":1,\"total_tokens\":4}}\n\ndata: [DONE]\n\n"; + +fn request(stream: bool) -> Value { + json!({ + "model": "public/model", + "messages": [{"role": "user", "content": "hi"}], + "stream": stream, + "stream_options": stream.then(|| json!({"include_usage": true})), + "caching": false, + "metadata": {"tag": "x"} + }) +} + +#[rstest] +#[case::json(false, COMPLETION, "application/json")] +#[case::sse(true, STREAM, "text/event-stream")] +#[tokio::test] +async fn openai_chat_passes_the_provider_response_through( + #[case] stream: bool, + #[case] upstream_body: &'static str, + #[case] content_type: &str, + #[values("/v1/chat/completions", "/engines/public%2Fmodel/chat/completions")] route: &str, +) { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(header("authorization", "Bearer test-key")) + .and(body_json(json!({ + "model": "gpt-6-astra", + "messages": [{"role": "user", "content": "hi"}], + "stream": stream, + "stream_options": stream.then(|| json!({"include_usage": true})), + }))) + .respond_with(ResponseTemplate::new(200).set_body_raw(upstream_body, content_type)) + .expect(1) + .mount(&upstream) + .await; + + let response = support::post( + support::app("openai/gpt-6-astra", &format!("{}/v1", upstream.uri())), + route, + request(stream), + ) + .await; + + assert_eq!(response.status(), 200); + assert_eq!(response.headers()["content-type"], content_type); + assert_eq!( + to_bytes(response.into_body(), 1 << 20).await.unwrap(), + upstream_body + ); +} + +#[rstest] +#[case::json(false)] +#[case::sse(true)] +#[tokio::test] +async fn an_openai_error_keeps_its_status_and_message(#[case] stream: bool) { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(429).set_body_raw( + r#"{"error":{"message":"Rate limit reached","type":"requests","code":"rate_limit_exceeded"}}"#, + "application/json", + )) + .mount(&upstream) + .await; + + let response = support::post( + support::app("openai/gpt-6-astra", &upstream.uri()), + "/v1/chat/completions", + request(stream), + ) + .await; + + assert_eq!(response.status(), 429); + let body = support::json(response).await; + assert_eq!(body["error"]["message"], "Rate limit reached"); + assert_eq!(body["error"]["type"], "rate_limit_error"); +} + +/// Sends the SSE head and one event, then holds the connection open until told to send the +/// rest, or drops it short of the end when `truncate` is set. +async fn paused_upstream(release: tokio::sync::oneshot::Receiver<()>, truncate: bool) -> String { + let listener = tokio::net::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; 64 * 1024]; + let _ = socket.read(&mut request).await; + socket + .write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n{:x}\r\n{FIRST_EVENT}\r\n", + FIRST_EVENT.len() + ) + .as_bytes(), + ) + .await + .unwrap(); + let _ = release.await; + if truncate { + return; + } + let _ = socket + .write_all(format!("{:x}\r\n{DONE}\r\n0\r\n\r\n", DONE.len()).as_bytes()) + .await; + }); + base +} + +const FIRST_EVENT: &str = "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"a\"}}]}\n\n"; +const DONE: &str = "data: [DONE]\n\n"; + +#[rstest] +#[tokio::test] +async fn the_gateway_forwards_each_event_while_the_provider_is_still_streaming() { + let (release, released) = tokio::sync::oneshot::channel(); + let base = paused_upstream(released, false).await; + let response = support::post( + support::app("openai/gpt-6-astra", &base), + "/v1/chat/completions", + request(true), + ) + .await; + assert_eq!(response.status(), 200); + let mut body = response.into_body().into_data_stream(); + + let first = tokio::time::timeout(Duration::from_secs(2), body.next()) + .await + .expect("the first event reaches the client before the provider finishes") + .unwrap() + .unwrap(); + assert_eq!(first, FIRST_EVENT); + release.send(()).unwrap(); + assert_eq!(body.next().await.unwrap().unwrap(), DONE); + assert!(body.next().await.is_none()); +} + +#[rstest] +#[tokio::test] +async fn a_stream_cut_by_the_provider_ends_with_an_openai_error_event() { + let (release, released) = tokio::sync::oneshot::channel(); + let base = paused_upstream(released, true).await; + let response = support::post( + support::app("openai/gpt-6-astra", &base), + "/v1/chat/completions", + request(true), + ) + .await; + release.send(()).unwrap(); + + let body = to_bytes(response.into_body(), 1 << 20).await.unwrap(); + let text = std::str::from_utf8(&body).unwrap(); + let error = text + .strip_prefix(FIRST_EVENT) + .and_then(|rest| rest.strip_prefix("data: ")) + .and_then(|rest| rest.strip_suffix("\n\n")) + .unwrap_or_else(|| panic!("expected the first event then an error event, got {text:?}")); + let error: Value = serde_json::from_str(error).unwrap(); + assert_eq!(error["error"]["code"], 500); +} + +#[rstest] +#[tokio::test] +async fn openai_like_chat_streaming_is_still_declined() { + let upstream = MockServer::start().await; + let response = support::post( + support::app("openai_like/gpt-6-astra", &upstream.uri()), + "/v1/chat/completions", + json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], "stream": true}), + ) + .await; + assert_eq!(response.status(), 501); + assert!(upstream.received_requests().await.unwrap().is_empty()); +} diff --git a/litellm-rust/crates/inference-chat/AGENTS.md b/litellm-rust/crates/inference-chat/AGENTS.md index a73880c2858..2e0d3846464 100644 --- a/litellm-rust/crates/inference-chat/AGENTS.md +++ b/litellm-rust/crates/inference-chat/AGENTS.md @@ -1,3 +1,4 @@ - Shared base and layering rules: [`../inference/AGENTS.md`](../inference/AGENTS.md) - This crate owns Chat Completions call orchestration. Payloads belong in `litellm-llms-types`, transformations in `llms/src/base_llm/chat` and `llms/src//chat` -- Chat Completions is currently non-streaming: the route returns its completed response, cached through `litellm_inference::caching::execute_unary` +- Translated providers (`ChatCompletions` in `src/route.rs`) are non-streaming: the route returns its completed response, cached through `litellm_inference::caching::execute_unary` +- OpenAI (`OpenAiChatCompletions` in `src/openai.rs`) is a passthrough machine: the provider's JSON body or SSE bytes go to the host as received, with no total timeout on a stream, only a per-read idle limit diff --git a/litellm-rust/crates/inference-chat/Cargo.toml b/litellm-rust/crates/inference-chat/Cargo.toml index 130786e788d..f4c3ee35c83 100644 --- a/litellm-rust/crates/inference-chat/Cargo.toml +++ b/litellm-rust/crates/inference-chat/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +bytes.workspace = true futures-util.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-cache-response.workspace = true @@ -16,7 +17,9 @@ litellm-inference.workspace = true litellm-llms.workspace = true litellm-llms-types.workspace = true litellm-secrets.workspace = true +reqwest.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } +tokio = { workspace = true, features = ["time"] } tracing.workspace = true [dev-dependencies] @@ -26,5 +29,5 @@ litellm-http = { workspace = true, features = ["test-support"] } litellm-inference-testing.workspace = true litellm-tracing.workspace = true rstest.workspace = true -tokio.workspace = true +tokio = { workspace = true, features = ["io-util", "sync"] } wiremock.workspace = true diff --git a/litellm-rust/crates/inference-chat/src/lib.rs b/litellm-rust/crates/inference-chat/src/lib.rs index 65e82af6a8d..f3187e67ab2 100644 --- a/litellm-rust/crates/inference-chat/src/lib.rs +++ b/litellm-rust/crates/inference-chat/src/lib.rs @@ -5,6 +5,7 @@ pub use litellm_inference::RouteError as Error; mod common_utils; pub mod constants; pub(crate) mod handler; +pub mod openai; mod prepare; use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse; use prepare::{prepare_provider_request, resolve_request}; diff --git a/litellm-rust/crates/inference-chat/src/openai.rs b/litellm-rust/crates/inference-chat/src/openai.rs new file mode 100644 index 00000000000..9fd389966fa --- /dev/null +++ b/litellm-rust/crates/inference-chat/src/openai.rs @@ -0,0 +1,610 @@ +use std::{convert::Infallible, time::Duration}; + +use bytes::Bytes; +use futures_util::{StreamExt, stream::BoxStream}; +use litellm_core_utils::get_llm_provider_logic::LlmProviders; +use litellm_host::{ + call::{CallOutput, HostedMachine, hosted_call}, + interceptors::{ + Cost, ExecutionFacts, Interceptors, ProviderIdentity, RequestContext, ResultSource, + WireRequest, + }, + observation::ObservationSender, + protocol::Protocol, +}; +use litellm_http::{outbound::OutboundRequest, request::truncate_error_body, transport}; +use litellm_inference::{CallOptions, context::CallContext, provider::resolve_llm_provider}; +use litellm_llms::{ + base_llm::auth::resolve_auth, + openai::chat::transformation::{OPENAI_CHAT_COMPLETIONS_CONFIG, is_stream}, +}; +use serde_json::{Map, Value}; + +use crate::{ + ChatCompletionsRoute, Error, common_utils::string_headers, + constants::CHAT_COMPLETIONS_TIMEOUT_SECS, +}; + +pub struct OpenAiChatCompletions; + +impl Protocol for OpenAiChatCompletions { + type Request = OpenAiChatCall; + /// The provider's JSON body, as it sent it. + type Response = Bytes; + type Error = Error; + type HostCall = Infallible; + type Chunk = Bytes; + type StreamHead = (); +} + +pub type OpenAiChatOutput = CallOutput; + +/// A Chat Completions body bound for an OpenAI deployment. `body` is the caller's request; +/// `model` is the deployment's, with or without its `openai/` prefix. +pub struct OpenAiChatCall { + pub model: String, + pub body: Map, + pub api_key: Option, + pub api_base: Option, + pub custom_llm_provider: Option, + pub extra_headers: Option>, + pub timeout: Option, +} + +impl ChatCompletionsRoute { + pub fn serves_openai(model: &str, custom_llm_provider: Option<&str>) -> bool { + resolve_llm_provider(model, custom_llm_provider, "chat completions") + .is_ok_and(|resolved| resolved.provider == LlmProviders::Openai) + } + + pub fn openai_machine( + self, + call: OpenAiChatCall, + observers: Option, + ) -> HostedMachine { + hosted_call( + call, + observers, + move |call, _, interceptors, observers| async move { + let context = CallContext::new( + &interceptors, + CallOptions { + cache: None, + observers, + }, + ); + Box::pin(self.run_openai(call, context)).await + }, + ) + } + + pub async fn execute_openai( + &self, + call: OpenAiChatCall, + interceptors: &impl Interceptors, + observers: Option, + ) -> Result { + let context = CallContext::new( + interceptors, + CallOptions { + cache: None, + observers, + }, + ); + litellm_host::lifecycle::observe_call( + context.observers.clone(), + Box::pin(self.run_openai(call, context)), + ) + .await + } + + #[tracing::instrument(name = "litellm.route", skip_all, fields( + route = "chat_completions", + model = %call.model, + provider, + resolved_model, + stream, + outcome + ))] + async fn run_openai( + &self, + call: OpenAiChatCall, + context: CallContext<'_, impl Interceptors>, + ) -> Result { + litellm_inference::diagnostic::call(async { + let resolved = resolve_llm_provider( + &call.model, + call.custom_llm_provider.as_deref(), + "chat completions", + )?; + if resolved.provider != LlmProviders::Openai { + return Err(Error::InvalidProvider( + <&str>::from(resolved.provider).to_string(), + )); + } + let model = resolved.model.to_string(); + let provider = <&str>::from(LlmProviders::Openai); + litellm_inference::diagnostic::provider(&model, provider); + let config = &OPENAI_CHAT_COMPLETIONS_CONFIG; + let secrets = self + .secrets + .resolve(config.secret_names(call.api_key.as_deref(), call.api_base.as_deref())) + .await?; + let lookup = |key: &str| secrets.get(key); + let body = config.transform_request(&model, call.body)?; + let environment = config.validate_environment( + string_headers(call.extra_headers)?, + call.api_key.as_deref(), + &lookup, + )?; + let url = config.get_complete_url(call.api_base.as_deref(), &lookup); + let authenticated = resolve_auth(&self.auth, environment, &lookup).await?; + let request_context = RequestContext { + model: model.clone(), + custom_llm_provider: provider.to_string(), + optional_params: Value::Object( + body.iter() + .filter(|(name, _)| !matches!(name.as_str(), "model" | "messages")) + .map(|(name, value)| (name.clone(), value.clone())) + .collect(), + ), + secret_fields: Vec::new(), + api_key: call.api_key.map(litellm_auth::SecretValue::new), + }; + let wire = context + .interceptors + .before_provider_request( + WireRequest { + url, + headers: authenticated.headers, + body: Value::Object(body), + }, + request_context, + ) + .await?; + let stream = wire + .body + .as_object() + .map(is_stream) + .transpose()? + .unwrap_or(false); + let timeout = call + .timeout + .unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS)); + let response = send(&self.http, wire, stream, timeout).await?; + let output = if stream { + CallOutput::Stream { + head: (), + chunks: chunks(response, timeout), + } + } else { + let body = response.bytes().await.map_err(network)?; + context + .response_received(&String::from_utf8_lossy(&body)) + .await?; + CallOutput::Complete(body) + }; + context + .result_ready(ExecutionFacts { + provider: ProviderIdentity { + model, + provider: provider.to_string(), + }, + source: ResultSource::Provider, + cost: Cost::Deferred, + }) + .await?; + Ok(output) + }) + .await + } +} + +/// A non-streaming call keeps the usual total timeout. A stream has no total limit, so a +/// 30 minute generation is not cut off: `timeout` bounds the wait for the response head +/// here and each read after it in [`chunks`]. +async fn send( + http: &litellm_http::Client, + wire: WireRequest, + stream: bool, + timeout: Duration, +) -> Result { + let request = OutboundRequest::json( + wire.url, + wire.headers, + &wire.body, + (!stream).then_some(timeout), + )?; + let response = tokio::time::timeout(timeout, litellm_inference::outbound::send(request, http)) + .await + .map_err(|_| silent_provider(timeout))? + .map_err(|error| Error::Transport(transport::Error::from_reqwest_before_dispatch(error)))?; + let status = response.status(); + if status.is_success() { + return Ok(response); + } + let text = response.text().await.map_err(network)?; + Err(Error::Transport(transport::Error::Http { + status: status.as_u16(), + body: truncate_error_body(&text), + })) +} + +fn chunks(response: reqwest::Response, idle: Duration) -> BoxStream<'static, Result> { + futures_util::stream::try_unfold(response, move |mut response| async move { + let chunk = tokio::time::timeout(idle, response.chunk()) + .await + .map_err(|_| silent_provider(idle))? + .map_err(network)?; + Ok(chunk.map(|chunk| (chunk, response))) + }) + .boxed() +} + +fn network(error: reqwest::Error) -> Error { + Error::Transport(transport::Error::from(error)) +} + +fn silent_provider(timeout: Duration) -> Error { + Error::Transport(transport::Error::Network(format!( + "the provider sent nothing for {}s", + timeout.as_secs_f64() + ))) +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use futures_util::TryStreamExt; + use litellm_host::interceptors::RawResponse; + use rstest::{fixture, rstest}; + use serde_json::json; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use wiremock::{ + Mock, MockServer, Request, ResponseTemplate, + matchers::{header, method, path}, + }; + + use super::*; + + const COMPLETION: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-6-astra","system_fingerprint":"fp","choices":[{"index":0,"message":{"role":"assistant","content":null,"tool_calls":[{"id":"call_1","type":"function","function":{"name":"f","arguments":"{}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}"#; + const STREAM: &str = "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\ndata: {\"id\":\"c\",\"choices\":[],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":1,\"total_tokens\":4}}\n\ndata: [DONE]\n\n"; + + #[derive(Default)] + struct Recording { + contexts: Mutex>, + raw: Mutex>, + results: Mutex>, + } + + impl Interceptors for Recording { + async fn before_provider_request( + &self, + wire: WireRequest, + context: RequestContext, + ) -> Result { + self.contexts.lock().unwrap().push(context); + Ok(wire) + } + + async fn after_provider_response(&self, raw: RawResponse) -> Result<(), Error> { + self.raw.lock().unwrap().push(raw.body); + Ok(()) + } + + async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { + self.results.lock().unwrap().push(facts); + Ok(()) + } + } + + #[fixture] + fn route() -> ChatCompletionsRoute { + ChatCompletionsRoute::new( + litellm_http::Client::plain_for_test(), + Arc::new(litellm_auth::AuthServices::default()), + litellm_inference_testing::no_secrets(), + ) + } + + fn call(api_base: &str, body: Value) -> OpenAiChatCall { + let Value::Object(body) = body else { + panic!("the body is an object"); + }; + OpenAiChatCall { + model: "openai/gpt-6-astra".into(), + body, + api_key: Some("sk-deployment".into()), + api_base: Some(format!("{api_base}/v1")), + custom_llm_provider: None, + extra_headers: None, + timeout: Some(Duration::from_secs(5)), + } + } + + async fn run( + route: &ChatCompletionsRoute, + call: OpenAiChatCall, + interceptors: &Recording, + ) -> Result { + route + .run_openai(call, CallContext::new(interceptors, CallOptions::default())) + .await + } + + async fn drain(chunks: BoxStream<'static, Result>) -> Result, Error> { + Ok(chunks.try_collect::>().await?.concat()) + } + + async fn collect(output: OpenAiChatOutput) -> Result, Error> { + match output { + CallOutput::Complete(body) => panic!("expected a stream, got {body:?}"), + CallOutput::Stream { chunks, .. } => drain(chunks).await, + } + } + + #[rstest] + #[case::openai_prefix("openai/gpt-6-astra", None, true)] + #[case::explicit_provider("gpt-6-astra", Some("openai"), true)] + #[case::openai_like("openai_like/gpt-6-astra", None, false)] + #[case::anthropic("anthropic/claude", None, false)] + #[case::unprefixed("gpt-6-astra", None, false)] + fn only_openai_deployments_take_the_passthrough( + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] expected: bool, + ) { + assert_eq!( + ChatCompletionsRoute::serves_openai(model, provider), + expected + ); + } + + #[rstest] + #[tokio::test] + async fn a_completion_returns_the_provider_body_byte_for_byte(route: ChatCompletionsRoute) { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(header("authorization", "Bearer sk-deployment")) + .respond_with(ResponseTemplate::new(200).set_body_raw(COMPLETION, "application/json")) + .expect(1) + .mount(&upstream) + .await; + let interceptors = Recording::default(); + + let output = run( + &route, + call( + &upstream.uri(), + json!({"model": "public", "messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "function", "function": {"name": "f"}}], "metadata": {"k": "v"}}), + ), + &interceptors, + ) + .await + .unwrap(); + + let CallOutput::Complete(body) = output else { + panic!("a non-streaming call completes"); + }; + assert_eq!(body, COMPLETION.as_bytes()); + let requests = upstream.received_requests().await.unwrap(); + let sent: Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert_eq!( + sent, + json!({"model": "gpt-6-astra", "messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "function", "function": {"name": "f"}}]}) + ); + assert_eq!(*interceptors.raw.lock().unwrap(), [COMPLETION]); + let contexts = interceptors.contexts.lock().unwrap(); + assert_eq!( + ( + contexts[0].model.as_str(), + contexts[0].custom_llm_provider.as_str() + ), + ("gpt-6-astra", "openai") + ); + assert_eq!( + contexts[0].optional_params, + json!({"tools": [{"type": "function", "function": {"name": "f"}}]}) + ); + assert_eq!(interceptors.results.lock().unwrap().len(), 1); + } + + #[rstest] + #[tokio::test] + async fn a_stream_relays_every_byte_including_the_usage_chunk(route: ChatCompletionsRoute) { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_raw(STREAM, "text/event-stream")) + .mount(&upstream) + .await; + let interceptors = Recording::default(); + + let output = run( + &route, + call( + &upstream.uri(), + json!({"model": "public", "messages": [{"role": "user", "content": "hi"}], + "stream": true, "stream_options": {"include_usage": true}}), + ), + &interceptors, + ) + .await + .unwrap(); + + assert_eq!(collect(output).await.unwrap(), STREAM.as_bytes()); + let request: &Request = &upstream.received_requests().await.unwrap()[0]; + let sent: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(sent["stream_options"], json!({"include_usage": true})); + assert!(interceptors.raw.lock().unwrap().is_empty()); + } + + /// Serves an SSE head and the first event, then waits for `release` before the rest. + async fn paused_upstream(release: tokio::sync::oneshot::Receiver<()>) -> String { + let listener = tokio::net::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; 64 * 1024]; + 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") + .await + .unwrap(); + let event = "data: {\"n\":1}\n\n"; + socket + .write_all(format!("{:x}\r\n{event}\r\n", event.len()).as_bytes()) + .await + .unwrap(); + let _ = release.await; + let done = "data: [DONE]\n\n"; + let _ = socket + .write_all(format!("{:x}\r\n{done}\r\n0\r\n\r\n", done.len()).as_bytes()) + .await; + }); + base + } + + #[rstest] + #[tokio::test] + async fn each_event_is_delivered_before_the_stream_finishes(route: ChatCompletionsRoute) { + let (release, released) = tokio::sync::oneshot::channel(); + let base = paused_upstream(released).await; + let output = run( + &route, + call(&base, json!({"model": "m", "messages": [], "stream": true})), + &Recording::default(), + ) + .await + .unwrap(); + let CallOutput::Stream { mut chunks, .. } = output else { + panic!("a streaming call streams"); + }; + + let first = tokio::time::timeout(Duration::from_secs(2), chunks.try_next()) + .await + .expect("the first event arrives while the provider is still streaming") + .unwrap() + .unwrap(); + assert_eq!(first, "data: {\"n\":1}\n\n"); + release.send(()).unwrap(); + assert_eq!(drain(chunks).await.unwrap(), b"data: [DONE]\n\n"); + } + + #[rstest] + #[tokio::test] + async fn a_stream_outlives_the_request_timeout_while_data_keeps_coming( + route: ChatCompletionsRoute, + ) { + let (release, released) = tokio::sync::oneshot::channel(); + let base = paused_upstream(released).await; + let mut streaming = call(&base, json!({"model": "m", "messages": [], "stream": true})); + streaming.timeout = Some(Duration::from_millis(300)); + let output = run(&route, streaming, &Recording::default()).await.unwrap(); + let CallOutput::Stream { mut chunks, .. } = output else { + panic!("a streaming call streams"); + }; + chunks.try_next().await.unwrap(); + + tokio::time::sleep(Duration::from_millis(200)).await; + release.send(()).unwrap(); + tokio::time::sleep(Duration::from_millis(200)).await; + + assert_eq!(drain(chunks).await.unwrap(), b"data: [DONE]\n\n"); + } + + #[rstest] + #[tokio::test] + async fn a_stream_that_goes_silent_past_the_timeout_fails(route: ChatCompletionsRoute) { + let (_release, released) = tokio::sync::oneshot::channel(); + let base = paused_upstream(released).await; + let mut streaming = call(&base, json!({"model": "m", "messages": [], "stream": true})); + streaming.timeout = Some(Duration::from_millis(100)); + let output = run(&route, streaming, &Recording::default()).await.unwrap(); + + assert!(matches!( + collect(output).await, + Err(Error::Transport(transport::Error::Network(_))) + )); + } + + #[rstest] + #[tokio::test] + async fn an_upstream_error_keeps_its_status(route: ChatCompletionsRoute) { + let upstream = MockServer::start().await; + let body = r#"{"error":{"message":"slow down","type":"rate_limit_error"}}"#; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(429).set_body_raw(body, "application/json")) + .mount(&upstream) + .await; + + let error = run( + &route, + call( + &upstream.uri(), + json!({"model": "m", "messages": [], "stream": true}), + ), + &Recording::default(), + ) + .await + .err() + .expect("the upstream failure fails the call"); + + assert_eq!( + error, + Error::Transport(transport::Error::Http { + status: 429, + body: body.into() + }) + ); + } + + #[rstest] + #[tokio::test] + async fn the_key_falls_back_to_the_environment(route: ChatCompletionsRoute) { + let upstream = MockServer::start().await; + Mock::given(header("authorization", "Bearer sk-env")) + .respond_with(ResponseTemplate::new(200).set_body_raw(COMPLETION, "application/json")) + .expect(1) + .mount(&upstream) + .await; + let route = ChatCompletionsRoute { + secrets: Arc::new(litellm_inference_testing::RecordingSecrets::new([( + "OPENAI_API_KEY", + "sk-env", + )])), + ..route + }; + let mut without_key = call(&upstream.uri(), json!({"model": "m", "messages": []})); + without_key.api_key = None; + + run(&route, without_key, &Recording::default()) + .await + .unwrap(); + } + + #[rstest] + #[tokio::test] + async fn a_missing_key_fails_before_any_request(route: ChatCompletionsRoute) { + let mut without_key = call("http://127.0.0.1:1", json!({"model": "m", "messages": []})); + without_key.api_key = None; + assert!(matches!( + run(&route, without_key, &Recording::default()).await, + Err(Error::Auth(litellm_auth::Error::MissingApiKey { .. })) + )); + } + + #[rstest] + #[tokio::test] + async fn a_non_openai_deployment_is_refused(route: ChatCompletionsRoute) { + let mut other = call("http://127.0.0.1:1", json!({"model": "m", "messages": []})); + other.model = "openai_like/gpt".into(); + assert_eq!( + run(&route, other, &Recording::default()).await.err(), + Some(Error::InvalidProvider("openai_like".into())) + ); + } +} diff --git a/litellm-rust/crates/inference-chat/tests/live/main.rs b/litellm-rust/crates/inference-chat/tests/live/main.rs new file mode 100644 index 00000000000..586d310caa7 --- /dev/null +++ b/litellm-rust/crates/inference-chat/tests/live/main.rs @@ -0,0 +1,4 @@ +//! Real provider calls, one module per provider. Each test skips when its provider's key is not +//! in the environment, so `cargo test` stays offline by default. + +mod openai; diff --git a/litellm-rust/crates/inference-chat/tests/live/openai.rs b/litellm-rust/crates/inference-chat/tests/live/openai.rs new file mode 100644 index 00000000000..f764b1b87fe --- /dev/null +++ b/litellm-rust/crates/inference-chat/tests/live/openai.rs @@ -0,0 +1,108 @@ +use std::time::Duration; + +use futures_util::TryStreamExt; +use litellm_host::call::CallOutput; +use litellm_inference_chat::{ + ChatCompletionsRoute, + openai::{OpenAiChatCall, OpenAiChatOutput}, +}; +use litellm_inference_testing::live::{LiveResources, within_deadline}; +use rstest::{fixture, rstest}; +use serde_json::{Map, Value, json}; + +const DEFAULT_MODEL: &str = "openai/gpt-6-astra"; + +/// A route reading `OPENAI_API_KEY` from the environment; tests skip when it is unset. +#[fixture] +fn route() -> Option { + std::env::var("OPENAI_API_KEY").ok()?; + let resources = LiveResources::default(); + Some(ChatCompletionsRoute::new( + resources.http, + resources.auth, + resources.secrets, + )) +} + +fn call(extra: Value) -> OpenAiChatCall { + let Value::Object(extra) = extra else { + panic!("extra params are an object"); + }; + let body: Map = json!({ + "model": "public", + "messages": [{"role": "user", "content": "Reply with the single word: pong"}], + "max_completion_tokens": 256, + }) + .as_object() + .unwrap() + .clone() + .into_iter() + .chain(extra) + .collect(); + OpenAiChatCall { + model: std::env::var("OPENAI_LIVE_MODEL").unwrap_or_else(|_| DEFAULT_MODEL.into()), + body, + api_key: None, + api_base: None, + custom_llm_provider: None, + extra_headers: None, + timeout: Some(Duration::from_secs(60)), + } +} + +async fn execute(route: &ChatCompletionsRoute, call: OpenAiChatCall) -> OpenAiChatOutput { + within_deadline(route.execute_openai(call, &(), None)) + .await + .unwrap() +} + +#[rstest] +#[tokio::test] +async fn a_completion_returns_openai_json(route: Option) { + let Some(route) = route else { return }; + let CallOutput::Complete(body) = execute(&route, call(json!({}))).await else { + panic!("expected a complete response"); + }; + let body: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["object"], "chat.completion", "{body}"); + assert!( + body["choices"][0]["message"]["content"] + .as_str() + .is_some_and(|text| text.to_ascii_lowercase().contains("pong")), + "{body}" + ); + assert!( + body["usage"]["total_tokens"].as_u64().unwrap() > 0, + "{body}" + ); +} + +#[rstest] +#[tokio::test] +async fn a_stream_ends_with_usage_and_done(route: Option) { + let Some(route) = route else { return }; + let output = execute( + &route, + call(json!({"stream": true, "stream_options": {"include_usage": true}})), + ) + .await; + let CallOutput::Stream { chunks, .. } = output else { + panic!("expected SSE"); + }; + let wire = within_deadline(chunks.try_collect::>()) + .await + .unwrap() + .concat(); + let text = std::str::from_utf8(&wire).unwrap(); + let events: Vec<&str> = text + .split("\n\n") + .filter_map(|event| event.trim().strip_prefix("data:")) + .map(str::trim) + .collect(); + assert_eq!(events.last(), Some(&"[DONE]"), "{text}"); + let usage: Value = serde_json::from_str(events[events.len() - 2]).unwrap(); + assert!( + usage["usage"]["total_tokens"].as_u64().unwrap() > 0, + "{usage}" + ); +} diff --git a/litellm-rust/crates/llms-types/src/formats/chat_completions/AGENTS.md b/litellm-rust/crates/llms-types/src/formats/chat_completions/AGENTS.md index 34c9b02c012..552a20a06a9 100644 --- a/litellm-rust/crates/llms-types/src/formats/chat_completions/AGENTS.md +++ b/litellm-rust/crates/llms-types/src/formats/chat_completions/AGENTS.md @@ -1,7 +1,7 @@ # rules - Chat Completions originated at OpenAI and is the most widely cloned LLM API. Types here describe the format, not OpenAI's deployment of it -- Spoken natively by OpenAI and OpenAI-compatible hosts, served by `llms/src/openai_like/chat/` +- Spoken natively by OpenAI, passed through by `llms/src/openai/chat/`, and by OpenAI-compatible hosts, served by `llms/src/openai_like/chat/` - Other providers translate into it from their own API (Anthropic Messages, Bedrock Converse) in `llms/src//chat/` - Compatible hosts commonly - accept a subset of parameters, which is recorded per provider in `supported_openai_param_mappings`, not by trimming the type diff --git a/litellm-rust/crates/llms/src/openai/chat/AGENTS.md b/litellm-rust/crates/llms/src/openai/chat/AGENTS.md new file mode 100644 index 00000000000..356742bc872 --- /dev/null +++ b/litellm-rust/crates/llms/src/openai/chat/AGENTS.md @@ -0,0 +1,6 @@ +# rules + +- Format spec lives in `llms-types/src/formats/chat_completions/AGENTS.md` +- Passthrough, not a `BaseConfig`: that contract parses messages and normalizes the response into `ChatCompletionsResponse`, which drops tool calls, ids and every field it does not type. OpenAI's own body and SSE bytes reach the caller as sent +- The request changes only by the deployment model name, the `extra_body` merge the OpenAI SDK does, and dropping the LiteLLM-owned keys Python consumes before the call +- Endpoint and credential fallbacks are shared with Responses in `../common_utils.rs` diff --git a/litellm-rust/crates/llms/src/openai/chat/mod.rs b/litellm-rust/crates/llms/src/openai/chat/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/llms/src/openai/chat/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/llms/src/openai/chat/transformation.rs b/litellm-rust/crates/llms/src/openai/chat/transformation.rs new file mode 100644 index 00000000000..43a3df53510 --- /dev/null +++ b/litellm-rust/crates/llms/src/openai/chat/transformation.rs @@ -0,0 +1,214 @@ +//! `litellm/llms/openai/chat/gpt_transformation.py` as a passthrough: the caller's body is +//! already OpenAI-shaped, so only the model name and LiteLLM's own request keys change, and +//! the response bytes reach the caller untouched. + +use serde_json::{Map, Value}; + +use crate::{ + Error, ErrorDetail, + base_llm::auth::{Headers, ValidatedEnvironment}, + openai::common_utils, +}; + +pub const OPENAI_CHAT_COMPLETIONS_PATH: &str = "/chat/completions"; + +/// The keys of Python's `all_litellm_params` a proxy caller can put in a chat request body. +/// Python consumes them before the OpenAI call, so they never reach the provider. +const LITELLM_OWNED_PARAMS: &[&str] = &[ + "api_base", + "api_key", + "api_version", + "base_model", + "cache", + "caching", + "context_window_fallback_dict", + "custom_llm_provider", + "extra_headers", + "fallbacks", + "guardrails", + "litellm_metadata", + "litellm_session_id", + "litellm_trace_id", + "max_retries", + "metadata", + "mock_response", + "mock_timeout", + "model_info", + "no-log", + "num_retries", + "prompt_id", + "prompt_variables", + "prompt_version", + "request_timeout", + "stream_timeout", + "tags", + "timeout", +]; +const INTERNAL_PARAM_PREFIX: &str = "_litellm_"; +const EXTRA_BODY: &str = "extra_body"; + +pub struct OpenAiChatConfig; + +pub const OPENAI_CHAT_COMPLETIONS_CONFIG: OpenAiChatConfig = OpenAiChatConfig; + +impl OpenAiChatConfig { + pub fn secret_names( + &self, + api_key: Option<&str>, + api_base: Option<&str>, + ) -> &'static [&'static str] { + common_utils::secret_names(api_key, api_base) + } + + pub fn validate_environment( + &self, + headers: Headers, + api_key: Option<&str>, + lookup: &dyn Fn(&str) -> Option, + ) -> Result { + common_utils::bearer_environment(headers, api_key, lookup) + } + + pub fn get_complete_url( + &self, + api_base: Option<&str>, + lookup: &dyn Fn(&str) -> Option, + ) -> String { + common_utils::complete_url(api_base, OPENAI_CHAT_COMPLETIONS_PATH, lookup) + } + + /// The OpenAI SDK merges `extra_body` over the body, and the deployment's model name + /// replaces the public one. + pub fn transform_request( + &self, + model: &str, + body: Map, + ) -> Result, Error> { + let mut body = body; + let extra_body = match body.remove(EXTRA_BODY) { + None | Some(Value::Null) => Map::new(), + Some(Value::Object(extra)) => extra, + Some(other) => { + return Err(Error::InvalidRequest(ErrorDetail::InvalidValue { + field: EXTRA_BODY, + expected: "an object", + actual: other, + })); + } + }; + let mut body: Map = body + .into_iter() + .chain(extra_body) + .filter(|(name, _)| !is_litellm_owned(name)) + .collect(); + body.insert("model".into(), Value::String(model.into())); + Ok(body) + } +} + +/// Whether the body asks for SSE. Anything but a boolean or null is the caller's mistake, +/// since the call cannot tell which response shape to expect. +pub fn is_stream(body: &Map) -> Result { + match body.get("stream") { + None | Some(Value::Null) => Ok(false), + Some(Value::Bool(stream)) => Ok(*stream), + Some(other) => Err(Error::InvalidRequest(ErrorDetail::InvalidValue { + field: "stream", + expected: "a boolean", + actual: other.clone(), + })), + } +} + +fn is_litellm_owned(name: &str) -> bool { + LITELLM_OWNED_PARAMS.contains(&name) || name.starts_with(INTERNAL_PARAM_PREFIX) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn object(value: Value) -> Map { + let Value::Object(map) = value else { + panic!("expected an object, got {value}"); + }; + map + } + + #[rstest] + fn the_body_passes_through_with_only_the_model_renamed() { + let body = json!({ + "model": "public/gpt", + "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "u"}}]}], + "tools": [{"type": "function", "function": {"name": "f", "parameters": {}}}], + "stream": true, + "stream_options": {"include_usage": true}, + "reasoning_effort": "high", + "a_future_param": {"nested": [1, 2]}, + "user": "end-user" + }); + let sent = OPENAI_CHAT_COMPLETIONS_CONFIG + .transform_request("gpt-6-astra", object(body.clone())) + .unwrap(); + let mut expected = object(body); + expected.insert("model".into(), json!("gpt-6-astra")); + assert_eq!(sent, expected); + } + + #[rstest] + #[case::metadata("metadata")] + #[case::mock_response("mock_response")] + #[case::fallbacks("fallbacks")] + #[case::credentials("api_key")] + #[case::timeout("timeout")] + #[case::no_log("no-log")] + #[case::internal("_litellm_control")] + fn litellm_owned_keys_never_reach_openai(#[case] name: &str) { + let body = object(json!({"model": "m", "messages": [], name: "x"})); + let sent = OPENAI_CHAT_COMPLETIONS_CONFIG + .transform_request("gpt", body) + .unwrap(); + assert_eq!(sent, object(json!({"model": "gpt", "messages": []}))); + } + + #[rstest] + fn extra_body_is_merged_over_the_body_but_not_over_the_model() { + let body = object(json!({ + "model": "m", + "temperature": 1, + "extra_body": {"temperature": 0, "top_k": 5, "model": "other", "metadata": {}} + })); + let sent = OPENAI_CHAT_COMPLETIONS_CONFIG + .transform_request("gpt", body) + .unwrap(); + assert_eq!( + sent, + object(json!({"model": "gpt", "temperature": 0, "top_k": 5})) + ); + } + + #[rstest] + fn a_non_object_extra_body_is_rejected() { + let body = object(json!({"model": "m", "extra_body": "x"})); + assert!(matches!( + OPENAI_CHAT_COMPLETIONS_CONFIG.transform_request("gpt", body), + Err(Error::InvalidRequest(ErrorDetail::InvalidValue { + field: "extra_body", + .. + })) + )); + } + + #[rstest] + #[case::absent(json!({}), Some(false))] + #[case::null(json!({"stream": null}), Some(false))] + #[case::off(json!({"stream": false}), Some(false))] + #[case::on(json!({"stream": true}), Some(true))] + #[case::string(json!({"stream": "true"}), None)] + fn the_stream_flag_must_be_a_boolean(#[case] body: Value, #[case] expected: Option) { + assert_eq!(is_stream(&object(body)).ok(), expected); + } +} diff --git a/litellm-rust/crates/llms/src/openai/common_utils.rs b/litellm-rust/crates/llms/src/openai/common_utils.rs new file mode 100644 index 00000000000..538ef0ff7f5 --- /dev/null +++ b/litellm-rust/crates/llms/src/openai/common_utils.rs @@ -0,0 +1,137 @@ +//! Endpoint and credential resolution every OpenAI API config shares, mirroring the +//! `api_base` / `api_key` fallbacks in `litellm/llms/openai/openai.py`. + +use litellm_auth::{CredentialPlacement, SecretValue}; + +use crate::{ + Error, + base_llm::auth::{AuthScheme, Headers, ValidatedEnvironment}, +}; + +pub const OPENAI_DEFAULT_API_BASE: &str = "https://api.openai.com/v1"; +pub const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; +const OPENAI_API_BASE_ENVS: [&str; 2] = ["OPENAI_BASE_URL", "OPENAI_API_BASE"]; + +/// The secrets a call has to resolve: only the ones its deployment leaves unset. +pub fn secret_names(api_key: Option<&str>, api_base: Option<&str>) -> &'static [&'static str] { + match ( + api_key.is_some_and(|key| !key.is_empty()), + api_base.is_some_and(|base| !base.is_empty()), + ) { + (true, true) => &[], + (true, false) => &["OPENAI_BASE_URL", "OPENAI_API_BASE"], + (false, true) => &["OPENAI_API_KEY"], + (false, false) => &["OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENAI_API_BASE"], + } +} + +/// The deployment's `api_base`, else `OPENAI_BASE_URL` or `OPENAI_API_BASE`, else OpenAI's +/// own endpoint, with `path` appended. +pub fn complete_url( + api_base: Option<&str>, + path: &str, + lookup: &dyn Fn(&str) -> Option, +) -> String { + let base = api_base + .filter(|base| !base.is_empty()) + .map(str::to_owned) + .or_else(|| OPENAI_API_BASE_ENVS.into_iter().find_map(lookup)) + .unwrap_or_else(|| OPENAI_DEFAULT_API_BASE.into()); + format!("{}{path}", base.trim_end_matches('/')) +} + +/// The deployment's key, else `OPENAI_API_KEY`, sent as a bearer token. The configured key +/// replaces any `authorization` the caller forwarded. +pub fn bearer_environment( + headers: Headers, + api_key: Option<&str>, + lookup: &dyn Fn(&str) -> Option, +) -> Result { + let key = api_key + .filter(|key| !key.is_empty()) + .map(str::to_owned) + .or_else(|| lookup(OPENAI_API_KEY_ENV)) + .filter(|key| !key.is_empty()) + .ok_or(litellm_auth::Error::MissingApiKey { + provider: "OpenAI", + environment_variable: OPENAI_API_KEY_ENV, + })?; + Ok(ValidatedEnvironment { + headers: litellm_http::request::with_default_headers( + headers, + &[("content-type", "application/json")], + ), + auth: AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: SecretValue::new(key), + }, + }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn env(name: &str) -> Option { + match name { + "OPENAI_API_BASE" => Some("http://env-base/v1/".into()), + "OPENAI_API_KEY" => Some("sk-env".into()), + _ => None, + } + } + + fn no_env(_: &str) -> Option { + None + } + + #[rstest] + #[case::default(None, no_env, "https://api.openai.com/v1/chat/completions")] + #[case::empty_deployment_base(Some(""), no_env, "https://api.openai.com/v1/chat/completions")] + #[case::environment(None, env, "http://env-base/v1/chat/completions")] + #[case::deployment_wins( + Some("http://mock/d/1/v1/"), + env, + "http://mock/d/1/v1/chat/completions" + )] + fn the_url_prefers_the_deployment_then_the_environment( + #[case] api_base: Option<&str>, + #[case] lookup: fn(&str) -> Option, + #[case] expected: &str, + ) { + assert_eq!( + complete_url(api_base, "/chat/completions", &lookup), + expected + ); + } + + #[rstest] + #[case::deployment(Some("sk-deployment"), no_env, "sk-deployment")] + #[case::deployment_wins(Some("sk-deployment"), env, "sk-deployment")] + #[case::environment(None, env, "sk-env")] + #[case::empty_deployment_key(Some(""), env, "sk-env")] + fn the_key_prefers_the_deployment_then_the_environment( + #[case] api_key: Option<&str>, + #[case] lookup: fn(&str) -> Option, + #[case] expected: &str, + ) { + let validated = bearer_environment(Vec::new(), api_key, &lookup).unwrap(); + assert!(matches!( + validated.auth, + AuthScheme::Credential { placement: CredentialPlacement::Bearer, secret } + if secret.expose() == expected + )); + } + + #[rstest] + fn a_missing_key_names_the_environment_variable() { + assert_eq!( + bearer_environment(Vec::new(), None, &no_env).unwrap_err(), + Error::Auth(litellm_auth::Error::MissingApiKey { + provider: "OpenAI", + environment_variable: "OPENAI_API_KEY", + }) + ); + } +} diff --git a/litellm-rust/crates/llms/src/openai/mod.rs b/litellm-rust/crates/llms/src/openai/mod.rs index b396b037bc5..e95fcecca93 100644 --- a/litellm-rust/crates/llms/src/openai/mod.rs +++ b/litellm-rust/crates/llms/src/openai/mod.rs @@ -1 +1,3 @@ +pub mod chat; +pub mod common_utils; pub mod responses; diff --git a/litellm-rust/crates/llms/src/openai/responses/transformation.rs b/litellm-rust/crates/llms/src/openai/responses/transformation.rs index 99cc3559514..ecbbe95fffe 100644 --- a/litellm-rust/crates/llms/src/openai/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openai/responses/transformation.rs @@ -3,20 +3,19 @@ use litellm_llms_types::formats::responses::{ }; use serde_json::{Map, Value}; -use litellm_auth::{CredentialPlacement, SecretValue}; - use crate::{ Error, base_llm::{ - auth::{AuthScheme, ValidatedEnvironment}, + auth::ValidatedEnvironment, responses::transformation::{ BaseResponsesApiConfig, ResponsesWebSocketProviderConfig, ResponsesWsTransformResult, enforce_model, }, }, + openai::common_utils::{self, OPENAI_DEFAULT_API_BASE}, }; -pub const OPENAI_RESPONSES_DEFAULT_API_BASE: &str = "https://api.openai.com/v1"; +pub const OPENAI_RESPONSES_DEFAULT_API_BASE: &str = OPENAI_DEFAULT_API_BASE; pub const OPENAI_RESPONSES_PATH: &str = "/responses"; pub struct OpenAiResponsesApiConfig; @@ -107,15 +106,7 @@ impl BaseResponsesApiConfig for OpenAiResponsesApiConfig { api_key: Option<&str>, api_base: Option<&str>, ) -> &'static [&'static str] { - match ( - api_key.is_some_and(|key| !key.is_empty()), - api_base.is_some_and(|base| !base.is_empty()), - ) { - (true, true) => &[], - (true, false) => &["OPENAI_BASE_URL", "OPENAI_API_BASE"], - (false, true) => &["OPENAI_API_KEY"], - (false, false) => &["OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENAI_API_BASE"], - } + common_utils::secret_names(api_key, api_base) } fn validate_environment( @@ -124,25 +115,7 @@ impl BaseResponsesApiConfig for OpenAiResponsesApiConfig { api_key: Option<&str>, lookup: &dyn Fn(&str) -> Option, ) -> Result { - let key = api_key - .filter(|key| !key.is_empty()) - .map(str::to_owned) - .or_else(|| lookup("OPENAI_API_KEY")) - .filter(|key| !key.is_empty()) - .ok_or(litellm_auth::Error::MissingApiKey { - provider: "OpenAI", - environment_variable: "OPENAI_API_KEY", - })?; - Ok(ValidatedEnvironment { - headers: litellm_http::request::with_default_headers( - headers, - &[("content-type", "application/json")], - ), - auth: AuthScheme::Credential { - placement: CredentialPlacement::Bearer, - secret: SecretValue::new(key), - }, - }) + common_utils::bearer_environment(headers, api_key, lookup) } fn get_complete_url( @@ -150,13 +123,7 @@ impl BaseResponsesApiConfig for OpenAiResponsesApiConfig { api_base: Option<&str>, lookup: &dyn Fn(&str) -> Option, ) -> String { - let base = api_base - .filter(|base| !base.is_empty()) - .map(str::to_owned) - .or_else(|| lookup("OPENAI_BASE_URL")) - .or_else(|| lookup("OPENAI_API_BASE")) - .unwrap_or_else(|| OPENAI_RESPONSES_DEFAULT_API_BASE.into()); - format!("{}/responses", base.trim_end_matches('/')) + common_utils::complete_url(api_base, OPENAI_RESPONSES_PATH, lookup) } fn transform_responses_api_request( diff --git a/litellm-rust/docs/customer-proxy/openai_chat-flow.svg b/litellm-rust/docs/customer-proxy/openai_chat-flow.svg new file mode 100644 index 00000000000..62f5ad6aab4 --- /dev/null +++ b/litellm-rust/docs/customer-proxy/openai_chat-flow.svg @@ -0,0 +1,76 @@ + + + + + + + + POST /v1/chat/completions to an openai/ deployment + + + + Client + OpenAI SDK body + + + + gateway-inference chat_completions::handle + virtual key auth, resolve deployment, + authorize model, caching::prepare + + + + + serves_openai(model, + custom_llm_provider)? + + + + + Existing ChatCompletions route (unchanged) + openai_like, anthropic, bedrock: non-stream, stream = 501 + + no + + + + inference-chat openai::run_openai (HostedMachine) + 1. api_key: deployment, else OPENAI_API_KEY (Bearer) + 2. api_base: deployment, else OPENAI_BASE_URL / OPENAI_API_BASE, + else https://api.openai.com/v1, + /chat/completions + 3. body: merge extra_body, drop LiteLLM-owned keys + (metadata, fallbacks, _litellm_*, ...), set model + 4. before_provider_request interceptor, then send + every other field, incl. stream_options, passes verbatim + + yes + + + + OpenAI (or api_base) + application/json or + text/event-stream + + + + + + stream: false + reqwest total timeout = deployment timeout (default 600s) + body read once, after_provider_response sees it + CallOutput::Complete(bytes) sent as application/json, + byte for byte + upstream 4xx/5xx: same status, OpenAI error envelope + carrying the provider's error.message + + + stream: true + no total timeout; the timeout bounds the wait for the + response head and each chunk read (idle limit) + CallOutput::Stream: each reqwest chunk goes to the + host-http Sse encoder as is, no buffering or re-framing + upstream cut or idle past the limit after headers: + a final data: {"error": ...} event, then close + + + diff --git a/litellm-rust/loadtest/README.md b/litellm-rust/loadtest/README.md index 2632adcd1db..023abec5fd9 100644 --- a/litellm-rust/loadtest/README.md +++ b/litellm-rust/loadtest/README.md @@ -30,7 +30,7 @@ Images are tagged `customer-proxy-` so every result maps to a commit. ```sh mkdir -p out docker run --rm -v "$PWD/out:/work/out" litellm-loadtest: genconfig -out out \ - -mocks http://mockllm:9000 -chat-provider openai_like -keys 100 + -mocks http://mockllm:9000 -keys 100 export GATEWAY_IMAGE=litellm-rust-gateway: TOOLS_IMAGE=litellm-loadtest: docker compose up -d gateway mockllm procmon docker compose run --rm loadgen loadgen -workload /workloads/smoke.json \ @@ -93,8 +93,6 @@ gateway samples and mock hit counts, and prints the summary. (`KEY=` for `cloud-run.sh`). Once keys are read from Postgres, drop it and set `create_postgres`. - `/v1/messages` auth: `-messages-auth x-api-key` sends the key the way the Anthropic SDK does. The gateway reads only `Authorization: Bearer` today. -- Chat: deployments use `openai_like/` because the gateway rejects `openai/`, and chat streaming - returns 501. - Bedrock: `/v1/messages` to a `bedrock/` deployment returns 501 ("Bedrock invoke messages request shaping" is unimplemented in Rust). mockllm already serves Bedrock's `/invoke` and `/invoke-with-response-stream` (AWS event-stream framing), so set `-bedrock-fraction 0.5` diff --git a/litellm-rust/loadtest/cmd/genconfig/main.go b/litellm-rust/loadtest/cmd/genconfig/main.go index 13c9b582fff..fc2a8b385db 100644 --- a/litellm-rust/loadtest/cmd/genconfig/main.go +++ b/litellm-rust/loadtest/cmd/genconfig/main.go @@ -31,7 +31,7 @@ func main() { chatGroups := flag.Int("chat-groups", 4, "chat model groups") messagesGroups := flag.Int("messages-groups", 4, "messages model groups") deployments := flag.Int("deployments", 4, "deployments per model group") - chatProvider := flag.String("chat-provider", "openai", "provider prefix for chat deployments (openai_like until native openai lands in Rust)") + chatProvider := flag.String("chat-provider", "openai", "provider prefix for chat deployments (openai is served natively, streaming included; openai_like is non-streaming only)") keys := flag.Int("keys", 1000, "virtual keys to generate") keySeed := flag.String("key-seed", "", "derive keys from this seed so every machine generates the same set (default: random keys)") masterKey := flag.String("master-key", "sk-loadtest-master", "gateway master key") diff --git a/litellm-rust/loadtest/compose.yaml b/litellm-rust/loadtest/compose.yaml index c57fd380894..b590d16c75c 100644 --- a/litellm-rust/loadtest/compose.yaml +++ b/litellm-rust/loadtest/compose.yaml @@ -2,7 +2,7 @@ # containers. Not a perf number (one host), but the same images and workloads as the cloud rig. # # docker run --rm -v "$PWD/out:/work/out" litellm-loadtest genconfig -out out \ -# -mocks http://mockllm:9000 -chat-provider openai_like +# -mocks http://mockllm:9000 # GATEWAY_IMAGE=litellm-rust-gateway: TOOLS_IMAGE=litellm-loadtest: docker compose up -d gateway mockllm # docker compose run --rm loadgen loadgen -workload /workloads/smoke.json \ # -target http://gateway:4000 -key sk-loadtest-master -out out/smoke.jsonl diff --git a/litellm-rust/loadtest/terraform/aws/variables.tf b/litellm-rust/loadtest/terraform/aws/variables.tf index 2e7eadcdc4b..e9bf6765af8 100644 --- a/litellm-rust/loadtest/terraform/aws/variables.tf +++ b/litellm-rust/loadtest/terraform/aws/variables.tf @@ -103,9 +103,9 @@ variable "deployments_per_group" { } variable "chat_provider" { - description = "Provider prefix for chat deployments: openai_like until native openai lands in the Rust gateway, then openai." + description = "Provider prefix for chat deployments: openai is served natively, streaming included. openai_like is non-streaming only." type = string - default = "openai_like" + default = "openai" } variable "key_count" {