scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-14 04:09:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 97

Merge remote-tracking branch 'origin/main' into litellm_test_d6496f

Browse source
This commit is contained in:
Yuneng Jiang 2026-10-10 00:24:42 -07:00
commit 1e439dad8d
No known key found for this signature in database
145 changed files with 11258 additions and 650 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

72
.github/workflows/test-unit.yml vendored
View file

@ -335,6 +335,7 @@ jobs:
tests/unit/llms/snowflake
tests/unit/llms/soniox
tests/unit/llms/stability
tests/unit/llms/strands_decider
tests/unit/llms/tavily
tests/unit/llms/tencent
tests/unit/llms/tinyfish
@ -1268,6 +1269,75 @@ jobs:
fi
echo "$chart: all $declared declared test suites ran"
done
openai-sdk-compat:
name: OpenAI ${{ matrix.openai-version }} compatibility
permissions:
contents: read
pull-requests: read
runs-on: ubuntu-latest
timeout-minutes: 20
strategy:
fail-fast: false
matrix:
openai-version: ["2.20.0", "3.0.0", "3.25.0"]
env:
LITELLM_LOCAL_MODEL_COST_MAP: "True"
PYTHONPATH: ${{ github.workspace }}
UV_PYTHON: "3.10"
OPENAI_VERSION: ${{ matrix.openai-version }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Detect relevant changes
id: changes
timeout-minutes: 2
uses: ./.github/actions/detect-changes
- name: Set up Python
if: steps.changes.outputs.decision != 'skip'
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: ${{ env.UV_PYTHON }}
- name: Set up uv
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group dev --group proxy-dev --extra proxy --no-install-workspace
uv pip install "openai==${OPENAI_VERSION:?}" --exclude-newer-package openai=2026-10-07T00:00:00Z
uv run --no-sync python -c 'import os, openai; assert openai.__version__ == os.environ["OPENAI_VERSION"], openai.__version__'
- name: Run OpenAI SDK compatibility tests
if: steps.changes.outputs.decision != 'skip'
run: |
uv run --no-sync pytest \
tests/unit/llms/openai/test_openai.py \
tests/unit/llms/openai/test_openai_common_utils.py \
tests/unit/llms/openai/test_openai_file_content_streaming.py \
tests/unit/test_openai_embedding_encoding_format_default.py \
tests/unit/llms/inception/test_inception_completion_transformation.py::test_inception_fim_targets_fim_endpoint \
tests/unit/llms/inception/test_inception_completion_transformation.py::test_inception_fim_async \
tests/unit/llms/azure/test_azure_exception_mapping.py \
tests/unit/test_completion_timeout_resolution.py \
tests/unit/types/test_router.py \
tests/unit/test_exception_header_preservation.py \
tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py::test_speech_binary_response_still_streaming_reports_the_bytes_downloaded_so_far \
--tb=short -q
- name: Run workload identity tests with SDK 3
if: steps.changes.outputs.decision != 'skip' && startsWith(matrix.openai-version, '3.')
run: uv run --no-sync pytest tests/unit/llms/openai/test_openai_workload_identity.py --tb=short -q
- name: Check consumer types with the installed SDK
if: steps.changes.outputs.decision != 'skip'
run: uv run --no-sync basedpyright --pythonpath .venv/bin/python --project tests/typing/pyrightconfig.json
coverage:
name: coverage
needs: unit
@ -1316,7 +1386,7 @@ jobs:
unit-passed:
name: unit passed
permissions: {}
needs: [rust-bridge, assert-shard-coverage, unit, ui-unit, docs, helm]
needs: [rust-bridge, assert-shard-coverage, unit, ui-unit, docs, helm, openai-sdk-compat]
if: always()
runs-on: ubuntu-latest
timeout-minutes: 2

56
ci_cd/cost_map_guard.py
View file

@ -14,8 +14,9 @@ import argparse
import json
import subprocess
import sys
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors
@ -25,6 +26,7 @@ BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH)
BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
EVALUATION_OUTPUT_PRICE_SOURCES: Final[Mapping[str, str]] = MappingProxyType({})
CostMap = dict[str, object]
@ -51,7 +53,31 @@ def _rendered_schema(cost_map: CostMap) -> str:
return str(error)
def _file_failures(head: Snapshot, head_map: CostMap) -> tuple[str, ...]:
def _evaluation_output_price_failure(
key: str, entry: object, evaluation_output_price_sources: Mapping[str, str]
) -> str | None:
if not isinstance(entry, dict) or entry.get("mode") != "evaluation":
return None
output_cost: Final = entry.get("output_cost_per_token")
if (
not isinstance(output_cost, (int, float))
or isinstance(output_cost, bool)
or output_cost <= 0
or key in evaluation_output_price_sources
):
return None
return (
f"{key}: mode evaluation rows bill input only; output_cost_per_token must be 0 unless the vendor "
"prices output tokens for this model (add the source to EVALUATION_OUTPUT_PRICE_SOURCES)"
)
def _file_failures(
head: Snapshot,
head_map: CostMap,
*,
evaluation_output_price_sources: Mapping[str, str] = EVALUATION_OUTPUT_PRICE_SOURCES,
) -> tuple[str, ...]:
schema_text: Final = _rendered_schema(head_map)
if not schema_text.startswith("{"):
return (schema_text,)
@ -75,6 +101,16 @@ def _file_failures(head: Snapshot, head_map: CostMap) -> tuple[str, ...]:
f"{COST_MAP_PATH} does not validate against its schema: {error}"
for error in validation_errors(head_map, json.loads(schema_text))[:20]
),
*(
failure
for key, entry in head_map.items()
if (
failure := _evaluation_output_price_failure(
key, entry, evaluation_output_price_sources
)
)
is not None
),
)
@ -121,13 +157,25 @@ def contract_for(bot: bool, changed_files: Sequence[str]) -> str:
return "human PR, file checks only" if touches_cost_map(changed_files) else "human PR, cost map untouched"
def guard_failures(base: Snapshot, head: Snapshot, changed_files: Sequence[str], bot: bool) -> tuple[str, ...]:
def guard_failures(
base: Snapshot,
head: Snapshot,
changed_files: Sequence[str],
bot: bool,
*,
evaluation_output_price_sources: Mapping[str, str] = EVALUATION_OUTPUT_PRICE_SOURCES,
) -> tuple[str, ...]:
if not bot and not touches_cost_map(changed_files):
return ()
head_map: Final = _parse_object(head.cost_map, COST_MAP_PATH)
if isinstance(head_map, str):
return (head_map,)
return (*_file_failures(head, head_map), *(_bot_failures(base, head_map, changed_files) if bot else ()))
return (
*_file_failures(
head, head_map, evaluation_output_price_sources=evaluation_output_price_sources
),
*(_bot_failures(base, head_map, changed_files) if bot else ()),
)
def _git(*args: str) -> str | None:

10
litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs
View file

@ -12,7 +12,7 @@ use pyo3::{
pub(super) struct ChatCompletionsPythonHost(pub InferenceHost);
pub(super) fn project(
host: &InferenceHost,
host: &mut InferenceHost,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> PyResult<ChatCompletionsCall> {
@ -38,7 +38,7 @@ impl PythonBinding for ChatCompletionsPythonHost {
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> Result<ChatCompletionsCall, InvokeError<Error>> {
let call = project(&self.0, py, arguments).map_err(InvokeError::Python)?;
let call = project(&mut self.0, py, arguments).map_err(InvokeError::Python)?;
if call
.optional_params
.get("stream")
@ -94,8 +94,10 @@ impl PythonHostCalls<ChatCompletions> for ChatCompletionsPythonHost {
}
impl PythonOwned for ChatCompletionsPythonHost {
fn close(&mut self, _: Python<'_>) {}
fn close(&mut self, _: Python<'_>) {
self.0.close();
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.0.request)
self.0.traverse(visit)
}
}

2
litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs
View file

@ -22,7 +22,7 @@ fn run_chat_completions(
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
let host = InferenceHost::new(
call.resolved()?.unbind(),
call.view()?,
"litellm.rust_bridge.chat_completions.route_host",
);
run_inference::<ChatCompletionsRoute, _>(

32
litellm-rust/crates/python-bridge/src/routes/inference.rs
View file

@ -5,15 +5,20 @@ use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache};
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_host_python::{PythonBinding, PythonHostCalls, from_py, present};
use litellm_host_python::{PythonBinding, PythonHostCalls, from_py};
use litellm_http::transport::Error as TransportError;
use litellm_inference::RouteError;
use litellm_secrets::source::SecretSource;
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
use pyo3::{
exceptions::PyValueError,
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
};
use serde::Serialize;
use serde_json::{Map, Value};
use super::NativeCall;
use super::{NativeCall, RequestView};
use crate::{
errors::{RustUpstreamError, route_error_to_pyerr},
marshal::{
@ -23,7 +28,7 @@ use crate::{
};
pub(super) struct InferenceHost {
pub request: Py<PyDict>,
view: RequestView,
module: &'static str,
}
@ -34,16 +39,17 @@ pub(super) struct ProjectedCall {
}
impl InferenceHost {
pub fn new(request: Py<PyDict>, module: &'static str) -> Self {
Self { request, module }
pub fn new(view: RequestView, module: &'static str) -> Self {
Self { view, module }
}
pub fn project(
&self,
&mut self,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
input: &str,
) -> PyResult<ProjectedCall> {
self.view.resolve(py, arguments)?;
let argument = |name: &str| self.argument(py, arguments, name);
let string = |name: &str| -> PyResult<Option<String>> {
argument(name)?.map(|value| value.extract()).transpose()
@ -93,7 +99,15 @@ impl InferenceHost {
arguments: &Bound<'py, PyDict>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
present(arguments, self.request.bind(py), name)
self.view.argument(py, arguments, name)
}
pub fn close(&mut self) {
self.view.close();
}
pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.view.traverse(visit)
}
pub fn parameters(
@ -130,7 +144,7 @@ impl InferenceHost {
let mapped = py
.import(self.module)?
.getattr("map_failure")?
.call1((native.value(py), self.request.bind(py)))?;
.call1((native.value(py), self.view.request(py)?))?;
Ok(PyErr::from_value(mapped))
}
}

34
litellm-rust/crates/python-bridge/src/routes/messages/host.rs
View file

@ -2,7 +2,7 @@ use crate::cache::{CacheCall, Cached, PythonCache, Selection};
use litellm_host_python::{PythonHostCalls, PythonOwned};
use bytes::Bytes;
use litellm_host_python::{InvokeError, PythonBinding, from_py, present, to_py};
use litellm_host_python::{InvokeError, PythonBinding, from_py, to_py};
use litellm_http::transport::Error as TransportError;
use litellm_inference_messages::{
Error, MessagesCall, MessagesSettings, MessagesShaping, litellm_params, messages_body,
@ -23,6 +23,7 @@ use crate::{
errors::{RustUpstreamError, route_error_to_pyerr},
marshal::{optional_timeout, project_optional_fields, public_response, python_timeout_seconds},
python_settings::missing_module,
routes::RequestView,
};
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.messages.route_host";
@ -118,14 +119,14 @@ fn native_error(py: Python<'_>, error: Error) -> PyResult<PyErr> {
/// The Python side of the Messages route: projects the prepared arguments and builds the
/// public response, chunks and exceptions.
pub(super) struct MessagesPythonHost {
request: Py<PyDict>,
view: RequestView,
cache: PythonCache,
}
impl MessagesPythonHost {
pub(super) fn new(request: Py<PyDict>, asynchronous: bool) -> Self {
pub(super) fn new(view: RequestView, asynchronous: bool) -> Self {
Self {
request,
view,
cache: PythonCache::new(asynchronous),
}
}
@ -135,8 +136,7 @@ impl MessagesPythonHost {
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> PyResult<Result<MessagesCall, Error>> {
let request = self.request.bind(py);
let argument = |name: &str| present(arguments, request, name);
let argument = |name: &str| self.view.argument(py, arguments, name);
let string = |name: &str| -> PyResult<Option<String>> {
argument(name)?.map(|value| value.extract()).transpose()
};
@ -182,9 +182,9 @@ impl MessagesPythonHost {
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> PyResult<Option<Map<String, Value>>> {
let request = self.request.bind(py);
let mapping = |name: &str| -> PyResult<Option<Map<String, Value>>> {
present(arguments, request, name)?
self.view
.argument(py, arguments, name)?
.map(|value| from_py(&value))
.transpose()
};
@ -199,7 +199,8 @@ impl MessagesPythonHost {
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> PyResult<Option<ProviderSpecificHeaders>> {
present(arguments, self.request.bind(py), "provider_specific_header")?
self.view
.argument(py, arguments, "provider_specific_header")?
.map(|value| from_py(&value))
.transpose()
}
@ -226,9 +227,9 @@ impl MessagesPythonHost {
}
fn provider(&self, py: Python<'_>) -> String {
self.request
.bind(py)
.get_item("custom_llm_provider")
self.view
.request(py)
.and_then(|request| request.get_item("custom_llm_provider"))
.ok()
.flatten()
.and_then(|value| value.extract::<Option<String>>().ok().flatten())
@ -242,7 +243,7 @@ impl MessagesPythonHost {
let mapped = py
.import(ROUTE_HOST_MODULE)
.and_then(|module| module.getattr("map_failure"))
.and_then(|map| map.call1((error.value(py), self.request.bind(py), self.provider(py))))
.and_then(|map| map.call1((error.value(py), self.view.request(py)?, self.provider(py))))
.and_then(|mapped| {
mapped
.extract::<Py<pyo3::exceptions::PyBaseException>>()
@ -264,6 +265,9 @@ impl PythonBinding for MessagesPythonHost {
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> Result<(MessagesCall, Selection), InvokeError<Error>> {
self.view
.resolve(py, arguments)
.map_err(InvokeError::Python)?;
let selection =
crate::cache::configure(&mut self.cache, py, arguments, "anthropic_messages")
.map_err(InvokeError::Python)?;
@ -341,16 +345,18 @@ impl PythonHostCalls<Cached<Messages>> for MessagesPythonHost {
impl PythonOwned for MessagesPythonHost {
fn close(&mut self, _: Python<'_>) {
self.view.close();
self.cache.close();
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.request)?;
self.view.traverse(visit)?;
self.cache.traverse(visit)
}
}
#[cfg(test)]
mod tests {
use litellm_host_python::present;
use rstest::rstest;
use serde_json::json;

2
litellm-rust/crates/python-bridge/src/routes/messages/mod.rs
View file

@ -44,7 +44,7 @@ fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyR
},
))
},
MessagesPythonHost::new(call.resolved()?.unbind(), asynchronous),
MessagesPythonHost::new(call.view()?, asynchronous),
hooks,
asynchronous,
)

101
litellm-rust/crates/python-bridge/src/routes/mod.rs
View file

@ -14,7 +14,9 @@ use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol}
use litellm_host_python::{
HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args,
};
use litellm_host_python::{missing_state, present};
use pyo3::{
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDict, PyMapping, PyTuple},
};
@ -36,6 +38,69 @@ impl<'py> NativeCall<'py> {
fn resolved(&self) -> PyResult<Bound<'py, PyDict>> {
effective_py_args(&self.base, &self.kwargs)
}
/// The view a route host reads its arguments from.
fn view(&self) -> PyResult<RequestView> {
RequestView::new(&self.base, &self.kwargs)
}
}
/// The keyword view a route host reads: the signature base under the caller's keywords,
/// laid again once the hooks have rewritten them so a keyword a hook deleted stays deleted
/// instead of falling back to the caller's value.
pub(crate) struct RequestView(Option<RequestDicts>);
struct RequestDicts {
base: Py<PyDict>,
request: Py<PyDict>,
}
impl RequestView {
pub(crate) fn new(base: &Bound<'_, PyDict>, kwargs: &Bound<'_, PyDict>) -> PyResult<Self> {
Ok(Self(Some(RequestDicts {
base: base.clone().unbind(),
request: effective_py_args(base, kwargs)?.unbind(),
})))
}
fn dicts(&self) -> PyResult<&RequestDicts> {
self.0.as_ref().ok_or_else(missing_state)
}
pub(crate) fn resolve(
&mut self,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> PyResult<()> {
let dicts = self.0.as_mut().ok_or_else(missing_state)?;
dicts.request = effective_py_args(dicts.base.bind(py), arguments)?.unbind();
Ok(())
}
pub(crate) fn request<'py>(&self, py: Python<'py>) -> PyResult<&Bound<'py, PyDict>> {
Ok(self.dicts()?.request.bind(py))
}
pub(crate) fn argument<'py>(
&self,
py: Python<'py>,
arguments: &Bound<'py, PyDict>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
present(arguments, self.request(py)?, name)
}
pub(crate) fn close(&mut self) {
self.0 = None;
}
pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
let Some(dicts) = &self.0 else {
return Ok(());
};
visit.call(&dicts.base)?;
visit.call(&dicts.request)
}
}
fn mapping_dict<'py>(value: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyDict>> {
@ -130,6 +195,42 @@ mod tests {
.unwrap()
}
fn dict<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyDict> {
py.eval(&std::ffi::CString::new(source).unwrap(), None, None)
.unwrap()
.cast_into::<PyDict>()
.unwrap()
}
#[rstest]
#[case::deleted_keyword_stays_deleted("{'system': None}", "{'system': 'caller'}", "{}", None)]
#[case::deleted_keyword_without_default("{}", "{'system': 'caller'}", "{}", None)]
#[case::rewritten_keyword_wins(
"{'system': None}",
"{'system': 'caller'}",
"{'system': 'hook'}",
Some("hook")
)]
#[case::signature_default_survives("{'system': 'default'}", "{}", "{}", Some("default"))]
fn request_view_reads_the_hook_rewritten_keywords_over_the_base(
#[case] base: &str,
#[case] kwargs: &str,
#[case] prepared: &str,
#[case] expected: Option<&str>,
) {
Python::initialize();
Python::attach(|py| {
let mut view = super::RequestView::new(&dict(py, base), &dict(py, kwargs)).unwrap();
let prepared = dict(py, prepared);
view.resolve(py, &prepared).unwrap();
let value = view
.argument(py, &prepared, "system")
.unwrap()
.map(|value| value.extract::<String>().unwrap());
assert_eq!(value.as_deref(), expected);
});
}
#[rstest]
#[case::sync("transcription")]
#[case::asynchronous("atranscription")]

19
litellm-rust/crates/python-bridge/src/routes/ocr/host.rs
View file

@ -15,7 +15,7 @@ use super::{
errors::to_pyerr as ocr_error_to_pyerr,
project::{OcrHostHandles, SETTINGS_ERROR_MARKER, project_request},
};
use crate::marshal::public_response;
use crate::{marshal::public_response, routes::RequestView};
enum OcrHostData {
Unprojected,
@ -27,14 +27,14 @@ enum OcrHostData {
/// document as it goes), acquires Azure AD tokens, and builds the public response and
/// exception.
pub(super) struct OcrPythonHost {
request: Py<PyDict>,
view: RequestView,
data: OcrHostData,
}
impl OcrPythonHost {
pub(super) fn new(request: Py<PyDict>) -> Self {
pub(super) fn new(view: RequestView) -> Self {
Self {
request,
view,
data: OcrHostData::Unprojected,
}
}
@ -58,7 +58,8 @@ impl OcrPythonHost {
let OcrHostData::Unprojected = self.data else {
return Err(missing_state());
};
let (request, handles) = project_request(self.request.bind(py), arguments)?;
self.view.resolve(py, arguments)?;
let (request, handles) = project_request(self.view.request(py)?, arguments)?;
let caller_token = handles.azure_ad_token_provider.is_some();
self.data = OcrHostData::Projected(Box::new(handles));
Ok(OcrCall {
@ -78,7 +79,7 @@ impl OcrPythonHost {
let mapped = py
.import("litellm.rust_bridge.ocr.route_host")
.and_then(|module| module.getattr("map_failure"))
.and_then(|map| map.call1((error.value(py), self.request.bind(py), provider)))
.and_then(|map| map.call1((error.value(py), self.view.request(py)?, provider)))
.and_then(|mapped| mapped.extract::<Py<PyBaseException>>().map_err(PyErr::from));
match mapped {
Ok(mapped) => PyErr::from_value(mapped.into_bound(py).into_any()),
@ -160,10 +161,11 @@ impl PythonHostCalls<Ocr> for OcrPythonHost {
impl PythonOwned for OcrPythonHost {
fn close(&mut self, _: Python<'_>) {
self.view.close();
self.data = OcrHostData::Released;
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.request)?;
self.view.traverse(visit)?;
if let OcrHostData::Projected(handles) = &self.data
&& let Some(provider) = &handles.azure_ad_token_provider
{
@ -218,7 +220,8 @@ del provider
.unwrap()
.cast_into::<PyDict>()
.unwrap();
let mut host = OcrPythonHost::new(PyDict::new(py).unbind());
let empty = PyDict::new(py);
let mut host = OcrPythonHost::new(RequestView::new(&empty, &empty).unwrap());
assert!(host.decode_request(py, &kwargs).unwrap().caller_token);
locals.del_item("kwargs").unwrap();
drop(kwargs);

2
litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs
View file

@ -44,7 +44,7 @@ fn run_ocr(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult
let route = litellm_inference_ocr::OcrRoute::new(client);
Ok(route.machine(request, None))
},
OcrPythonHost::new(call.resolved()?.unbind()),
OcrPythonHost::new(call.view()?),
hooks,
asynchronous,
)

10
litellm-rust/crates/python-bridge/src/routes/responses/host.rs
View file

@ -12,7 +12,7 @@ use pyo3::{
pub(super) struct ResponsesPythonHost(pub InferenceHost);
pub(super) fn project(
host: &InferenceHost,
host: &mut InferenceHost,
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> PyResult<ResponsesCall> {
@ -59,7 +59,7 @@ impl PythonBinding for ResponsesPythonHost {
py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> Result<ResponsesCall, InvokeError<Error>> {
let call = project(&self.0, py, arguments).map_err(InvokeError::Python)?;
let call = project(&mut self.0, py, arguments).map_err(InvokeError::Python)?;
if call
.optional_params
.get("stream")
@ -121,8 +121,10 @@ impl PythonHostCalls<Responses> for ResponsesPythonHost {
}
impl PythonOwned for ResponsesPythonHost {
fn close(&mut self, _: Python<'_>) {}
fn close(&mut self, _: Python<'_>) {
self.0.close();
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.0.request)
self.0.traverse(visit)
}
}

2
litellm-rust/crates/python-bridge/src/routes/responses/mod.rs
View file

@ -61,7 +61,7 @@ fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> Py
"native Python responses streaming",
));
}
let host = InferenceHost::new(resolved.unbind(), ROUTE_HOST_MODULE);
let host = InferenceHost::new(call.view()?, ROUTE_HOST_MODULE);
run_inference::<ResponsesRoute, _>(py, call, asynchronous, ResponsesPythonHost(host))
}

5
litellm/__init__.py
View file

@ -100,6 +100,7 @@ from litellm.constants import (
DEFAULT_ALLOWED_FAILS,
)
import httpx
from openai import DefaultAsyncHttpxClient, DefaultHttpxClient
# register_async_client_cleanup is lazy-loaded and called on first access
@ -420,8 +421,8 @@ error_logs: Dict = {}
add_function_to_prompt: bool = (
False # if function calling not supported by api, append function call details to system prompt
)
client_session: Optional[httpx.Client] = None
aclient_session: Optional[httpx.AsyncClient] = None
client_session: Optional[Union[httpx.Client, DefaultHttpxClient]] = None
aclient_session: Optional[Union[httpx.AsyncClient, DefaultAsyncHttpxClient]] = None
model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks'
model_cost_map_url: str = os.getenv(
"LITELLM_MODEL_COST_MAP_URL",

70
litellm/decisions/main.py
View file

@ -1,13 +1,13 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from dataclasses import dataclass, field, replace
from typing import Final, TypeAlias
import httpx
from pydantic import TypeAdapter, ValidationError
from pydantic import ConfigDict, TypeAdapter, ValidationError
from typing_extensions import assert_never
import litellm
from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER
from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER, normalize_drop_params
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.decisions.transformation import (
@ -40,6 +40,9 @@ DecisionsRequestFormat: TypeAlias = DecisionsRequestBody | OpenAIDecisionRequest
_SYSTEMONE_REQUEST_ADAPTER: Final[TypeAdapter[DecisionsRequestBody]] = TypeAdapter(DecisionsRequestBody)
_OPENAI_REQUEST_ADAPTER: Final[TypeAdapter[OpenAIDecisionRequestBody]] = TypeAdapter(OpenAIDecisionRequestBody)
_SAFETY_IDENTIFIER_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(
str | None, config=ConfigDict(title="safety_identifier")
)
_HANDLER: Final = BaseLLMHTTPHandler()
@ -57,6 +60,7 @@ class _DecisionsCall:
logging_obj: LiteLLMLoggingObj | None
headers: Mapping[str, str]
timeout: float | httpx.Timeout | None
litellm_params: Mapping[str, object] = field(repr=False)
def _supported_providers() -> tuple[str, ...]:
@ -96,16 +100,46 @@ def _validate_request(
)
def _ir_request(request: DecisionsRequestFormat) -> DecisionsIRRequest:
def _ir_request(request: DecisionsRequestFormat, safety_identifier: str | None) -> DecisionsIRRequest:
match request:
case DecisionsRequestBody():
return systemone_request_to_ir(request)
return replace(systemone_request_to_ir(request), safety_identifier=safety_identifier)
case OpenAIDecisionRequestBody():
return openai_request_to_ir(request)
case _:
assert_never(request)
def _drops_params(kwargs: Mapping[str, object]) -> bool:
return litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True
def _request_safety_identifier(safety_identifier: object, kwargs: Mapping[str, object]) -> str | None:
if isinstance(safety_identifier, str) or not _drops_params(kwargs):
return _SAFETY_IDENTIFIER_ADAPTER.validate_python(safety_identifier)
return None
@dataclass(frozen=True, slots=True)
class _UnsupportedSafetyIdentifier:
pass
def _provider_ir_request(
request: DecisionsRequestFormat,
*,
safety_identifier: str | None,
provider_config: BaseDecisionsConfig,
kwargs: Mapping[str, object],
) -> DecisionsIRRequest | _UnsupportedSafetyIdentifier:
ir_request: Final = _ir_request(request, safety_identifier)
if ir_request.safety_identifier is None or provider_config.supports_safety_identifier:
return ir_request
if _drops_params(kwargs):
return replace(ir_request, safety_identifier=None)
return _UnsupportedSafetyIdentifier()
def _prepare_call(
*,
model: str,
@ -144,8 +178,12 @@ def _prepare_call(
except ValueError as error:
raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error
try:
request_safety_identifier: Final = _request_safety_identifier(safety_identifier, kwargs)
request: Final = _validate_request(
state=state, questions=questions, decision_input=decision_input, safety_identifier=safety_identifier
state=state,
questions=questions,
decision_input=decision_input,
safety_identifier=request_safety_identifier,
)
except ValidationError as error:
raise litellm.BadRequestError(
@ -169,7 +207,22 @@ def _prepare_call(
llm_provider=provider,
)
ir_request: Final = _ir_request(request)
ir_request: Final = _provider_ir_request(
request,
safety_identifier=request_safety_identifier,
provider_config=provider_config,
kwargs=kwargs,
)
if isinstance(ir_request, _UnsupportedSafetyIdentifier):
raise litellm.UnsupportedParamsError(
message=(
f"{provider} does not support parameters: ['safety_identifier'], for model={model}. "
"To drop these, set `litellm.drop_params=True` or for proxy:\n\n"
"`litellm_settings:\n drop_params: true`\n"
),
model=model,
llm_provider=provider,
)
body: Final = provider_config.transform_decisions_request(
model=canonical_model, request=ir_request, custom_llm_provider=provider
)
@ -204,6 +257,7 @@ def _prepare_call(
logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
headers=extra_headers or {},
timeout=timeout,
litellm_params=kwargs,
)
@ -287,6 +341,7 @@ async def adecisions(
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
litellm_params=call.litellm_params,
)
except Exception as error:
raise _map_upstream_exception(error, call) from error
@ -332,6 +387,7 @@ def decisions(
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
litellm_params=call.litellm_params,
)
except Exception as error:
raise _map_upstream_exception(error, call) from error

62
litellm/exceptions.py
View file

@ -241,14 +241,7 @@ class BadRequestError(openai.BadRequestError):
self.litellm_debug_info = litellm_debug_info
self.max_retries = max_retries
self.num_retries = num_retries
# Use response if it's a valid httpx.Response with a request, otherwise use minimal error response
# Note: We check _request (not .request property) to avoid RuntimeError when _request is None
if (
response is not None
and isinstance(response, httpx.Response)
and hasattr(response, "_request")
and getattr(response, "_request", None) is not None
):
if response is not None and getattr(response, "_request", None) is not None:
self.response = response
else:
self.response = _get_minimal_error_response()
@ -360,6 +353,8 @@ class UnprocessableEntityError(openai.UnprocessableEntityError):
class Timeout(openai.APITimeoutError):
request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request
def __init__(
self,
message,
@ -458,6 +453,9 @@ class RateLimitError(openai.RateLimitError):
:class:`RateLimitErrorCategory` for the available values.
"""
response: httpx.Response # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX response
request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request
def __init__(
self,
message,
@ -515,7 +513,9 @@ class RateLimitError(openai.RateLimitError):
),
)
super().__init__(
self.message, response=self.response, body=body
self.message,
response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime
body=body,
) # Call the base class constructor with the parameters it needs
self.code = "429"
self.type = "throttling_error"
@ -594,12 +594,7 @@ class PaymentRequiredError(BadRequestError):
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
) -> None:
response_is_valid: Final = (
response is not None
and isinstance(response, httpx.Response)
and hasattr(response, "_request")
and getattr(response, "_request", None) is not None
)
response_is_valid: Final = response is not None and getattr(response, "_request", None) is not None
response_for_parent: Final = (
response
if response_is_valid
@ -770,6 +765,9 @@ class ServiceUnavailableError(openai.APIStatusError):
class BadGatewayError(openai.APIStatusError):
response: httpx.Response # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX response
request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request
def __init__(
self,
message,
@ -797,7 +795,9 @@ class BadGatewayError(openai.APIStatusError):
),
)
super().__init__(
self.message, response=self.response, body=None
self.message,
response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime
body=None,
) # Call the base class constructor with the parameters it needs
def __str__(self):
@ -818,6 +818,9 @@ class BadGatewayError(openai.APIStatusError):
class InternalServerError(openai.InternalServerError):
response: httpx.Response # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX response
request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request
def __init__(
self,
message,
@ -846,7 +849,9 @@ class InternalServerError(openai.InternalServerError):
),
)
super().__init__(
self.message, response=self.response, body=body
self.message,
response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime
body=body,
) # Call the base class constructor with the parameters it needs
self.type = "internal_server_error"
@ -911,6 +916,8 @@ class APIError(openai.APIError):
# raised if an invalid request (not get, delete, put, post) is made
class APIConnectionError(openai.APIConnectionError):
request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request
def __init__(
self,
message,
@ -929,7 +936,10 @@ class APIConnectionError(openai.APIConnectionError):
self.request = httpx.Request(method="POST", url="https://api.openai.com/v1")
self.max_retries = max_retries
self.num_retries = num_retries
super().__init__(message=self.message, request=self.request)
super().__init__(
message=self.message,
request=self.request, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime
)
def __str__(self):
_message = self.message
@ -950,6 +960,9 @@ class APIConnectionError(openai.APIConnectionError):
# raised if an invalid request (not get, delete, put, post) is made
class APIResponseValidationError(openai.APIResponseValidationError):
response: httpx.Response # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX response
request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request
def __init__(
self,
message,
@ -967,7 +980,11 @@ class APIResponseValidationError(openai.APIResponseValidationError):
self.litellm_debug_info = litellm_debug_info
self.max_retries = max_retries
self.num_retries = num_retries
super().__init__(response=response, body=None, message=message)
super().__init__(
response=response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime
body=None,
message=message,
)
def __str__(self):
_message = self.message
@ -1087,6 +1104,9 @@ class BudgetExceededError(Exception):
## DEPRECATED ##
class InvalidRequestError(openai.BadRequestError):
response: httpx.Response # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX response
request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request
def __init__(self, message, model, llm_provider):
self.status_code = 400
self.message = message
@ -1097,7 +1117,9 @@ class InvalidRequestError(openai.BadRequestError):
request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object
)
super().__init__(
message=self.message, response=self.response, body=None
message=self.message,
response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime
body=None,
) # Call the base class constructor with the parameters it needs

6
litellm/integrations/code_interpreter_interception/handler.py
View file

@ -26,9 +26,11 @@ from litellm.types.integrations.code_interpreter_interception import (
)
from litellm.types.integrations.custom_logger import (
CHAT_COMPLETION_AGENTIC_SURFACE,
CODE_INTERPRETER_INTERCEPTION_PREFIX,
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
AgenticLoopPlan,
AgenticLoopRequestPatch,
as_converted_stream,
is_interception_internal_key,
)
from litellm.types.llms.openai import (
@ -290,8 +292,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
else:
kwargs[_SANDBOX_KEY] = uuid.uuid4().hex
if kwargs.get("stream"):
kwargs["stream"] = False
kwargs[_CONVERTED_STREAM_KEY] = True
kwargs.update(as_converted_stream(kwargs, CODE_INTERPRETER_INTERCEPTION_PREFIX))
kwargs.pop("stream_options", None)
self._write_interception_metadata(kwargs)
function_tool: Final = self._get_function_tool(call_type=call_type)

15
litellm/integrations/websearch_interception/handler.py
View file

@ -40,8 +40,12 @@ from litellm.llms.base_llm.search.transformation import SearchResponse
from litellm.types.integrations.custom_logger import (
CHAT_COMPLETION_AGENTIC_SURFACE,
RESPONSES_AGENTIC_SURFACE,
WEBSEARCH_CONVERTED_STREAM_KEY,
WEBSEARCH_INTERCEPTION_PREFIX,
WEBSEARCH_STREAM_OPTIONS_KEY,
AgenticLoopPlan,
AgenticLoopRequestPatch,
as_converted_stream,
)
from litellm.types.integrations.websearch_interception import (
AnthropicSearchQuery,
@ -489,8 +493,7 @@ class WebSearchInterceptionLogger(CustomLogger):
if kwargs.get("stream"):
verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False")
kwargs["stream"] = False
kwargs["_websearch_interception_converted_stream"] = True
return as_converted_stream(kwargs, WEBSEARCH_INTERCEPTION_PREFIX)
return kwargs
@ -511,8 +514,7 @@ class WebSearchInterceptionLogger(CustomLogger):
if kwargs.get("stream"):
verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False")
converted_kwargs["stream"] = False
converted_kwargs["_websearch_interception_converted_stream"] = True
return as_converted_stream(converted_kwargs, WEBSEARCH_INTERCEPTION_PREFIX)
return converted_kwargs
@ -670,7 +672,10 @@ class WebSearchInterceptionLogger(CustomLogger):
if kwargs.get("stream"):
verbose_logger.debug("WebSearchInterception: Converting stream=True to stream=False")
kwargs["stream"] = False
kwargs["_websearch_interception_converted_stream"] = True
kwargs[WEBSEARCH_CONVERTED_STREAM_KEY] = True
if "stream_options" in kwargs:
kwargs[WEBSEARCH_STREAM_OPTIONS_KEY] = kwargs["stream_options"]
del kwargs["stream_options"]
return kwargs

68
litellm/litellm_core_utils/chat_completion_agentic_loop.py
View file

@ -4,7 +4,9 @@ import json
from collections.abc import Mapping
from itertools import chain
from types import MappingProxyType
from typing import Final, cast
from typing import Final
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
@ -17,11 +19,12 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.types.integrations.custom_logger import (
CHAT_COMPLETION_AGENTIC_SURFACE,
HEADROOM_CONVERTED_STREAM_KEY,
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
AgenticLoopPlan,
AgenticLoopRequestPatch,
converted_stream_requested,
is_interception_internal_key,
stashed_stream_options,
)
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
@ -57,10 +60,14 @@ def _post_hook_overridden(callback: CustomLogger) -> bool:
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
def _converted_stream_requested(kwargs: Mapping[str, object]) -> bool:
return bool(
kwargs.get("_code_interpreter_interception_converted_stream") or kwargs.get(HEADROOM_CONVERTED_STREAM_KEY)
)
_STREAM_OPTIONS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
def _stashed_stream_options(kwargs: Mapping[str, object]) -> Mapping[str, object] | None:
try:
return _STREAM_OPTIONS_ADAPTER.validate_python(stashed_stream_options(kwargs))
except ValidationError:
return None
def _coerce_int(value: object, default: int) -> int:
@ -100,23 +107,28 @@ def _check_agentic_loop_safety(
return fingerprint
def _wrap_response_as_fake_stream(
def _converted_stream_replay(
response: object,
*,
kwargs: Mapping[str, object],
depth: int,
model: str,
custom_llm_provider: str,
logging_obj: object,
) -> object:
) -> CustomStreamWrapper | None:
if depth or not converted_stream_requested(kwargs):
return None
if isinstance(response, CustomStreamWrapper):
return response
if not isinstance(response, ModelResponse) or not isinstance(logging_obj, LiteLLMLoggingObject):
return response
return None
return CustomStreamWrapper(
completion_stream=MockResponseIterator(model_response=response),
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
stream_options=_stashed_stream_options(kwargs),
)
@ -213,14 +225,15 @@ async def _execute_chat_completion_agentic_plan(
model,
str(e),
)
if _converted_stream_requested(kwargs) and not depth:
return _wrap_response_as_fake_stream(
response_followup,
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
return response_followup
replay: Final = _converted_stream_replay(
response_followup,
kwargs=kwargs,
depth=depth,
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
return response_followup if replay is None else replay
finally:
try:
await callback.async_agentic_loop_cleanup_hook(plan=plan, kwargs=kwargs)
@ -240,7 +253,7 @@ async def maybe_run_chat_completion_agentic_loop(
model: str,
messages: list,
optional_params: dict,
kwargs: dict,
kwargs: dict[str, object],
logging_obj: object,
custom_llm_provider: str,
stream: bool,
@ -343,14 +356,11 @@ async def maybe_run_chat_completion_agentic_loop(
str(e),
)
if _converted_stream_requested(kwargs) and not depth:
return cast(
"ModelResponse | CustomStreamWrapper",
_wrap_response_as_fake_stream(
response,
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
),
)
return None
return _converted_stream_replay(
response,
kwargs=kwargs,
depth=depth,
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)

25
litellm/litellm_core_utils/completion_timeout.py
View file

@ -6,6 +6,7 @@ from collections.abc import Callable
from typing import Final
import httpx
from openai import Timeout as SDKTimeout
from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS
@ -13,6 +14,12 @@ from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS
class CompletionTimeout:
"""Resolves HTTP timeout for ``completion()`` from model vs global settings."""
@staticmethod
def normalize(timeout: httpx.Timeout | SDKTimeout) -> httpx.Timeout:
if isinstance(timeout, httpx.Timeout): # pyright: ignore[reportUnnecessaryIsInstance] # SDK 2 aliases both classes.
return timeout
return httpx.Timeout(connect=timeout.connect, read=timeout.read, write=timeout.write, pool=timeout.pool)
@staticmethod
def _fallback_when_no_explicit_timeout(
global_timeout: float | str | None,
@ -31,7 +38,7 @@ class CompletionTimeout:
@staticmethod
def resolve(
model_timeout: float | str | httpx.Timeout | None,
model_timeout: float | str | httpx.Timeout | SDKTimeout | None,
kwargs: dict,
custom_llm_provider: str,
*,
@ -49,7 +56,7 @@ class CompletionTimeout:
Coerce :class:`httpx.Timeout` when the provider does not support it.
"""
resolved: float | str | httpx.Timeout
resolved: float | str | httpx.Timeout | SDKTimeout
if model_timeout is not None:
resolved = model_timeout
elif kwargs.get("timeout") is not None:
@ -59,12 +66,12 @@ class CompletionTimeout:
else:
resolved = CompletionTimeout._fallback_when_no_explicit_timeout(global_timeout)
if isinstance(resolved, httpx.Timeout) and not supports_httpx_timeout(custom_llm_provider):
read_timeout: Final = resolved.read
resolved = (
if isinstance(resolved, (httpx.Timeout, SDKTimeout)):
normalized: Final = CompletionTimeout.normalize(resolved)
if supports_httpx_timeout(custom_llm_provider):
return normalized
read_timeout: Final = normalized.read
return (
float(read_timeout) if read_timeout is not None else COMPLETION_HTTP_FALLBACK_SECONDS
) # default 10 min timeout
elif not isinstance(resolved, httpx.Timeout):
resolved = float(resolved)
return resolved
return float(resolved)

4
litellm/litellm_core_utils/realtime_streaming.py
View file

@ -1031,7 +1031,7 @@ class RealTimeStreaming:
cast(str, transcript),
item_id=cast(str | None, event.get("item_id")),
)
if not blocked and not self._is_transcription_session:
if not blocked and self._should_disable_vad_auto_response():
await self._send_to_backend(json.dumps({"type": "response.create"}))
continue
## LOGGING
@ -1078,7 +1078,7 @@ class RealTimeStreaming:
transcript,
item_id=event_obj.get("item_id"),
)
if not blocked and not self._is_transcription_session:
if not blocked and self._should_disable_vad_auto_response():
await self._send_to_backend(json.dumps({"type": "response.create"}))
return True
return False

3
litellm/llms/anthropic/pass_through/adapters/handler.py
View file

@ -24,6 +24,7 @@ from litellm.llms.anthropic.pass_through.utils import (
litellm_logging_obj_from_kwargs,
local_model_name,
)
from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -564,7 +565,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
# Maintainability: when adding a new Anthropic-only request param to
# ``AnthropicMessagesRequestOptionalParams``, also extend
# ``ANTHROPIC_ONLY_REQUEST_KEYS`` here so it doesn't silently leak.
excluded_keys: Final = ANTHROPIC_ONLY_REQUEST_KEYS | {"anthropic_messages"}
excluded_keys: Final = ANTHROPIC_ONLY_REQUEST_KEYS | CONVERTED_STREAM_KEYS | {"anthropic_messages"}
# NOTE: extra_kwargs was already coerced from None to {} at the top of
# this method (line ~220). It is guaranteed to be a dict here.
for key, value in extra_kwargs.items():

3
litellm/llms/anthropic/pass_through/responses_adapters/handler.py
View file

@ -8,6 +8,7 @@ from collections.abc import AsyncIterator, Coroutine, Mapping
from typing import Final, TypeAlias
import litellm
from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS
from litellm.types.llms.anthropic import (
AllAnthropicMessageValues,
AllAnthropicToolsValues,
@ -126,7 +127,7 @@ def _build_responses_kwargs(
responses_kwargs["stream"] = True
# Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.)
excluded: Final = frozenset(("anthropic_messages",))
excluded: Final = CONVERTED_STREAM_KEYS | {"anthropic_messages"}
for key, value in forwarded_kwargs.items():
if key == "litellm_logging_obj" and value is not None:
from litellm.litellm_core_utils.litellm_logging import (

8
litellm/llms/azure/common_utils.py
View file

@ -16,7 +16,7 @@ from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
from litellm.constants import DEFAULT_MAX_RETRIES
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.openai.common_utils import BaseOpenAILLM
from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient
from litellm.secret_managers.get_azure_ad_token_provider import (
get_azure_ad_token_provider,
)
@ -733,7 +733,11 @@ class BaseAzureLLM(BaseOpenAILLM):
azure_client_params: Final[_AzureGatewayClientParams] = {
"api_version": api_version,
"base_url": f"{api_base}",
"http_client": litellm.client_session,
"http_client": (
litellm.aclient_session or OpenAIAsyncHTTPClient()
if acompletion
else litellm.client_session or OpenAIHTTPClient()
),
"max_retries": max_retries,
"timeout": timeout,
}

15
litellm/llms/base_llm/decisions/transformation.py
View file

@ -301,6 +301,7 @@ class BaseDecisionsConfig(ABC):
api_key_env: tuple[str, ...] = ()
api_base_env: tuple[str, ...] = ()
api_key_required: bool = True
supports_safety_identifier: bool = False
health_check_questions: Mapping[str, Mapping[str, object]] = MappingProxyType(
{"reachable": MappingProxyType({"type": "noul", "instructions": "Is the service reachable?"})}
)
@ -340,6 +341,20 @@ class BaseDecisionsConfig(ABC):
def get_complete_url(self, api_base: str, model: str) -> str:
return f"{api_base.rstrip('/').removesuffix('/v1')}{self.path}"
def signs_with_aws(self, api_base: str) -> bool:
return False
def sign_request(
self,
headers: Mapping[str, str],
url: str,
api_base: str,
body: Mapping[str, object],
api_key: str | None,
litellm_params: Mapping[str, object],
) -> tuple[Mapping[str, str], bytes | None]:
return headers, None
def transform_decisions_request(
self,
model: str,

42
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -4,7 +4,7 @@ import json
import ssl
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence
from contextlib import asynccontextmanager
from functools import lru_cache
from functools import lru_cache, partial
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -1517,6 +1517,7 @@ class BaseLLMHTTPHandler:
api_key: str | None,
headers: Mapping[str, str],
timeout: float | httpx.Timeout | None,
litellm_params: Mapping[str, object],
client: HTTPHandler | None = None,
) -> DecisionsIRResponse:
url, outbound_headers, data = self._prepare_decisions_request(
@ -1528,10 +1529,23 @@ class BaseLLMHTTPHandler:
api_key=api_key,
headers=headers,
)
signed_headers, signed_body = provider_config.sign_request(
headers=outbound_headers,
url=url,
api_base=api_base,
body=data,
api_key=api_key,
litellm_params=litellm_params,
)
sync_httpx_client: Final = client if client is not None else get_httpx_client()
try:
response: Final = sync_httpx_client.post(
url, json=data, headers=outbound_headers, timeout=timeout, logging_obj=logging_obj
url,
content=signed_body,
json=data if signed_body is None else None,
headers=dict(signed_headers),
timeout=timeout,
logging_obj=logging_obj,
)
except httpx.HTTPError as e:
raise self._handle_error(e=e, provider_config=provider_config)
@ -1551,6 +1565,7 @@ class BaseLLMHTTPHandler:
api_key: str | None,
headers: Mapping[str, str],
timeout: float | httpx.Timeout | None,
litellm_params: Mapping[str, object],
client: AsyncHTTPHandler | None = None,
) -> DecisionsIRResponse:
url, outbound_headers, data = self._prepare_decisions_request(
@ -1562,6 +1577,18 @@ class BaseLLMHTTPHandler:
api_key=api_key,
headers=headers,
)
sign: Final = partial(
provider_config.sign_request,
headers=outbound_headers,
url=url,
api_base=api_base,
body=data,
api_key=api_key,
litellm_params=litellm_params,
)
signed_headers, signed_body = (
await run_aws_signing(sign) if provider_config.signs_with_aws(api_base) else sign()
)
async_httpx_client: Final = (
client
if client is not None
@ -1569,7 +1596,12 @@ class BaseLLMHTTPHandler:
)
try:
response: Final = await async_httpx_client.post(
url, json=data, headers=outbound_headers, timeout=timeout, logging_obj=logging_obj
url,
content=signed_body,
json=data if signed_body is None else None,
headers=dict(signed_headers),
timeout=timeout,
logging_obj=logging_obj,
)
except httpx.HTTPError as e:
raise self._handle_error(e=e, provider_config=provider_config)
@ -2495,12 +2527,12 @@ class BaseLLMHTTPHandler:
optional_param_names: Final = _responses_api_optional_request_param_names()
updated_response_params: Final = {
**response_api_optional_request_params,
**{key: value for key, value in response_api_optional_request_params.items() if key in modified_kwargs},
**{key: value for key, value in modified_kwargs.items() if key in optional_param_names},
}
updated_litellm_params: Final = GenericLiteLLMParams(
**{
**dict(litellm_params),
**{key: value for key, value in dict(litellm_params).items() if key in modified_kwargs},
**{
key: value
for key, value in modified_kwargs.items()

49
litellm/llms/openai/common_utils.py
View file

@ -2,6 +2,7 @@
Common helpers / utils across al OpenAI endpoints
"""
import asyncio
import hashlib
import inspect
import json
@ -10,12 +11,13 @@ import ssl
import time
import uuid
from collections.abc import AsyncIterator, Iterator, Mapping
from contextlib import suppress
from typing import TYPE_CHECKING, Final, Literal, NamedTuple, Optional
from urllib.parse import urlsplit
import httpx
import openai
from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, DefaultAsyncHttpxClient, DefaultHttpxClient, OpenAI
from openai.types.chat import ChatCompletion, ChatCompletionChunk, ChatCompletionMessage
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
@ -26,6 +28,7 @@ if TYPE_CHECKING:
from aiohttp import ClientSession
import litellm
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
from litellm.litellm_core_utils.token_counter import token_counter
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.custom_httpx.http_handler import (
@ -48,6 +51,42 @@ _AZURE_OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(Azur
_OPENAI_API_HOST: Final[str] = "api.openai.com"
_OPENAI_HTTPX_DEFAULT_TIMEOUT: Final = CompletionTimeout.normalize(openai.DEFAULT_TIMEOUT)
_OPENAI_HTTPX_CONNECTION_LIMITS: Final = httpx.Limits(
max_connections=openai.DEFAULT_CONNECTION_LIMITS.max_connections,
max_keepalive_connections=openai.DEFAULT_CONNECTION_LIMITS.max_keepalive_connections,
keepalive_expiry=openai.DEFAULT_CONNECTION_LIMITS.keepalive_expiry,
)
class OpenAIHTTPClient(httpx.Client):
def __init__(self) -> None:
super().__init__(
timeout=_OPENAI_HTTPX_DEFAULT_TIMEOUT,
limits=_OPENAI_HTTPX_CONNECTION_LIMITS,
follow_redirects=True,
)
def __del__(self) -> None:
with suppress(Exception):
if not self.is_closed:
self.close()
class OpenAIAsyncHTTPClient(httpx.AsyncClient):
def __init__(self) -> None:
super().__init__(
timeout=_OPENAI_HTTPX_DEFAULT_TIMEOUT,
limits=_OPENAI_HTTPX_CONNECTION_LIMITS,
follow_redirects=True,
)
def __del__(self) -> None:
with suppress(Exception):
if not self.is_closed:
asyncio.get_running_loop().create_task(self.aclose())
def is_openai_backed_api_base(api_base: str) -> bool:
hostname: Final = urlsplit(api_base).hostname
return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}"))
@ -221,7 +260,9 @@ class BaseOpenAILLM:
return _cached_client
@staticmethod
def owns_wrapped_http_client(http_client: httpx.Client | httpx.AsyncClient | None) -> bool:
def owns_wrapped_http_client(
http_client: httpx.Client | httpx.AsyncClient | DefaultHttpxClient | DefaultAsyncHttpxClient | None,
) -> bool:
"""Whether litellm may close an SDK client built around ``http_client``.
``_get_async_http_client`` / ``_get_sync_http_client`` hand back
@ -304,7 +345,7 @@ class BaseOpenAILLM:
@staticmethod
def _get_async_http_client(
shared_session: Optional["ClientSession"] = None,
) -> httpx.AsyncClient | None:
) -> httpx.AsyncClient | DefaultAsyncHttpxClient | None:
if litellm.aclient_session is not None:
return litellm.aclient_session
@ -337,7 +378,7 @@ class BaseOpenAILLM:
return cls._get_async_http_client(shared_session)
@staticmethod
def _get_sync_http_client() -> httpx.Client | None:
def _get_sync_http_client() -> httpx.Client | DefaultHttpxClient | None:
if litellm.client_session is not None:
return litellm.client_session

8
litellm/llms/openai/completion/handler.py
View file

@ -12,7 +12,7 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser
from litellm.types.utils import LlmProviders, ModelResponse, TextCompletionResponse
from litellm.utils import ProviderConfigManager
from ..common_utils import BaseOpenAILLM, OpenAIError
from ..common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIError, OpenAIHTTPClient
from .transformation import OpenAITextCompletionConfig
@ -129,7 +129,7 @@ class OpenAITextCompletion(BaseLLM):
openai_client = OpenAI(
api_key=api_key,
base_url=api_base,
http_client=litellm.client_session,
http_client=litellm.client_session or OpenAIHTTPClient(),
timeout=timeout,
max_retries=max_retries,
organization=organization,
@ -233,7 +233,7 @@ class OpenAITextCompletion(BaseLLM):
openai_client = OpenAI(
api_key=api_key,
base_url=api_base,
http_client=litellm.client_session,
http_client=litellm.client_session or OpenAIHTTPClient(),
timeout=timeout,
max_retries=max_retries,
organization=organization,
@ -290,7 +290,7 @@ class OpenAITextCompletion(BaseLLM):
openai_client = AsyncOpenAI(
api_key=api_key,
base_url=api_base,
http_client=litellm.aclient_session,
http_client=litellm.aclient_session or OpenAIAsyncHTTPClient(),
timeout=timeout,
max_retries=max_retries,
organization=organization,

1
litellm/llms/openai/decisions/transformation.py
View file

@ -292,6 +292,7 @@ def ir_to_openai_response(
class OpenAIDecisionsConfig(BaseDecisionsConfig):
path = "/v1/decisions"
supports_safety_identifier = True
def get_default_api_base(self) -> str | None:
return "https://api.openai.com"

5
litellm/llms/openai/fine_tuning/handler.py
View file

@ -6,6 +6,7 @@ from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
from openai.types.fine_tuning import FineTuningJob
from litellm._logging import verbose_logger
from litellm.llms.openai.common_utils import OpenAIAsyncHTTPClient, OpenAIHTTPClient
from litellm.types.utils import LiteLLMFineTuningJob
_AZURE_STATUS_MAP: Final[Mapping[object, str]] = {
@ -87,9 +88,9 @@ class OpenAIFineTuningAPI:
elif v is not None:
data[k] = v
if _is_async is True:
openai_client = AsyncOpenAI(**data)
openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient())
else:
openai_client = OpenAI(**data)
openai_client = OpenAI(**data, http_client=OpenAIHTTPClient())
else:
openai_client = client

24
litellm/llms/openai/image_variations/handler.py
View file

@ -14,7 +14,7 @@ from litellm.utils import ProviderConfigManager
from ...base_llm.image_variations.transformation import BaseImageVariationConfig
from ...custom_httpx.llm_http_handler import LiteLLMLoggingObj
from ..common_utils import OpenAIError
from ..common_utils import OpenAIAsyncHTTPClient, OpenAIError, OpenAIHTTPClient
class OpenAIImageVariationsHandler:
@ -23,22 +23,14 @@ class OpenAIImageVariationsHandler:
client: OpenAI | None,
init_client_params: dict,
):
if client is None:
openai_client = OpenAI(
**init_client_params,
)
else:
openai_client = client
return openai_client
if client is not None:
return client
return OpenAI(**init_client_params, http_client=litellm.client_session or OpenAIHTTPClient())
def get_async_client(self, client: AsyncOpenAI | None, init_client_params: dict) -> AsyncOpenAI:
if client is None:
openai_client = AsyncOpenAI(
**init_client_params,
)
else:
openai_client = client
return openai_client
if client is not None:
return client
return AsyncOpenAI(**init_client_params, http_client=litellm.aclient_session or OpenAIAsyncHTTPClient())
async def async_image_variations(
self,
@ -62,7 +54,6 @@ class OpenAIImageVariationsHandler:
init_client_params: Final = {
"api_key": api_key,
"base_url": api_base,
"http_client": litellm.client_session,
"timeout": timeout,
"max_retries": max_retries,
"organization": organization,
@ -179,7 +170,6 @@ class OpenAIImageVariationsHandler:
init_client_params: Final = {
"api_key": api_key,
"base_url": api_base,
"http_client": litellm.client_session,
"timeout": timeout,
"max_retries": max_retries,
"organization": organization,

14
litellm/llms/openai/openai.py
View file

@ -53,7 +53,9 @@ from .chat.gpt_transformation import OpenAIGPTConfig, OpenAIUnknownModelConfig
from .chat.o_series_transformation import OpenAIOSeriesConfig
from .common_utils import (
BaseOpenAILLM,
OpenAIAsyncHTTPClient,
OpenAIError,
OpenAIHTTPClient,
build_output_token_limit_response,
drop_params_from_unprocessable_entity_error,
is_openai_backed_api_base,
@ -1749,9 +1751,9 @@ class OpenAIFilesAPI(BaseLLM):
elif v is not None:
data[k] = v
if _is_async is True:
openai_client = AsyncOpenAI(**data)
openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient())
else:
openai_client = OpenAI(**data)
openai_client = OpenAI(**data, http_client=OpenAIHTTPClient())
else:
openai_client = client
@ -2107,9 +2109,9 @@ class OpenAIBatchesAPI(BaseLLM):
elif v is not None:
data[k] = v
if _is_async is True:
openai_client = AsyncOpenAI(**data)
openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient())
else:
openai_client = OpenAI(**data)
openai_client = OpenAI(**data, http_client=OpenAIHTTPClient())
else:
openai_client = client
@ -2317,7 +2319,7 @@ class OpenAIAssistantsAPI(BaseLLM):
data["base_url"] = v
elif v is not None:
data[k] = v
openai_client = OpenAI(**data)
openai_client = OpenAI(**data, http_client=OpenAIHTTPClient())
else:
openai_client = client
@ -2342,7 +2344,7 @@ class OpenAIAssistantsAPI(BaseLLM):
data["base_url"] = v
elif v is not None:
data[k] = v
openai_client = AsyncOpenAI(**data)
openai_client = AsyncOpenAI(**data, http_client=OpenAIAsyncHTTPClient())
else:
openai_client = client

154
litellm/llms/strands_decider/decisions/transformation.py
View file

@ -1,7 +1,161 @@
import hashlib
import json
import re
from collections.abc import Mapping
from dataclasses import dataclass
from functools import partial
from types import MappingProxyType
from typing import Final
from urllib.parse import quote, unquote
import httpx
from pydantic import BaseModel, ConfigDict, ValidationError
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, sign_aws_json_post
from litellm.types.decisions import DecisionsIRRequest, DecisionsIRResponse
from litellm.types.llms.bedrock import AwsAuthParams
AGENTCORE_SESSION_HEADER: Final = "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"
_RESERVED_HEADER_PREFIXES: Final = ("x-amzn-bedrock-agentcore-runtime-", "x-amz-")
_RUNTIME_ARN: Final = re.compile(
r"\Aarn:aws(?:-[a-z]+)*:bedrock-agentcore:(?P<region>[a-z0-9-]+):[0-9]{12}:runtime/[A-Za-z0-9_-]+\Z"
)
_INVOCATIONS_PATH: Final = re.compile(r"\A/runtimes/(?P<arn>[^/]+)/invocations\Z")
_RUNTIME_ERROR_STATUS: Final = MappingProxyType({"bad_request": 400, "invalid_request": 400, "loading": 503})
@dataclass(frozen=True, slots=True)
class AgentCoreRuntime:
arn: str
region: str
@property
def invocations_url(self) -> str:
host: Final = f"bedrock-agentcore.{self.region}.{get_aws_dns_suffix(self.region)}"
return f"https://{host}/runtimes/{quote(self.arn, safe='')}/invocations"
@property
def default_session_id(self) -> str:
return f"litellm-decider-{hashlib.sha256(self.arn.encode()).hexdigest()}"
def agentcore_runtime(api_base: str) -> AgentCoreRuntime | None:
match: Final = _RUNTIME_ARN.match(api_base)
return None if match is None else AgentCoreRuntime(arn=api_base, region=match["region"])
def invoked_runtime(url: httpx.URL) -> AgentCoreRuntime | None:
match: Final = _INVOCATIONS_PATH.match(url.raw_path.decode("ascii"))
runtime: Final = None if match is None else agentcore_runtime(unquote(match["arn"]))
return runtime if runtime is not None and runtime.invocations_url == str(url) else None
class _RuntimeError(BaseModel):
model_config = ConfigDict(frozen=True)
code: str
message: str = ""
class _RuntimeReply(BaseModel):
model_config = ConfigDict(frozen=True)
error: _RuntimeError | None = None
answers: object = None
def _runtime_error(content: bytes) -> _RuntimeError | None:
try:
reply: Final = _RuntimeReply.model_validate_json(content)
except ValidationError:
return None
return reply.error if reply.answers is None else None
class _SessionParams(BaseModel):
model_config = ConfigDict(frozen=True)
agentcore_runtime_session_id: str | None = None
def _is_reserved_header(name: str) -> bool:
lowered: Final = name.lower()
return lowered == "host" or lowered.startswith(_RESERVED_HEADER_PREFIXES)
def _runtime_headers(
headers: Mapping[str, str], runtime: AgentCoreRuntime, litellm_params: Mapping[str, object]
) -> Mapping[str, str]:
dropped: Final = sorted(name for name in headers if _is_reserved_header(name))
if dropped:
verbose_logger.warning("Strands Decider: not forwarding reserved AgentCore header(s) %s", dropped)
configured: Final = _SessionParams.model_validate(litellm_params).agentcore_runtime_session_id
return {
**{name: value for name, value in headers.items() if not _is_reserved_header(name)},
AGENTCORE_SESSION_HEADER: runtime.default_session_id if configured is None else configured,
}
class StrandsDeciderDecisionsConfig(BaseDecisionsConfig):
api_key_env = ("STRANDS_DECIDER_API_KEY",)
api_base_env = ("STRANDS_DECIDER_API_BASE",)
api_key_required = False
def __init__(self, aws: BaseAWSLLM | None = None) -> None:
self._aws: Final = aws if aws is not None else BaseAWSLLM()
def get_complete_url(self, api_base: str, model: str) -> str:
runtime: Final = agentcore_runtime(api_base)
return super().get_complete_url(api_base, model) if runtime is None else runtime.invocations_url
def signs_with_aws(self, api_base: str) -> bool:
return agentcore_runtime(api_base) is not None
def sign_request(
self,
headers: Mapping[str, str],
url: str,
api_base: str,
body: Mapping[str, object],
api_key: str | None,
litellm_params: Mapping[str, object],
) -> tuple[Mapping[str, str], bytes | None]:
runtime: Final = agentcore_runtime(api_base)
if runtime is None:
return headers, None
session_headers: Final = _runtime_headers(headers, runtime, litellm_params)
if api_key is not None:
return session_headers, None
payload: Final = json.dumps(body)
signed: Final = sign_aws_json_post(
get_credentials=partial(
self._aws.resolve_credentials, AwsAuthParams.model_validate(litellm_params), runtime.region
),
service_name="bedrock-agentcore",
aws_region_name=runtime.region,
url=url,
body=payload,
headers=session_headers,
)
return dict(signed.headers.items()), payload.encode()
def transform_decisions_response(
self,
model: str,
custom_llm_provider: str,
raw_response: httpx.Response,
request: DecisionsIRRequest,
) -> DecisionsIRResponse:
runtime_error: Final = (
None if invoked_runtime(raw_response.request.url) is None else _runtime_error(raw_response.content)
)
if runtime_error is not None:
raise BaseLLMException(
status_code=_RUNTIME_ERROR_STATUS.get(runtime_error.code, 500),
message=f"Strands Decider runtime error '{runtime_error.code}': {runtime_error.message}",
)
return super().transform_decisions_response(model, custom_llm_provider, raw_response, request)

42
litellm/main.py
View file

@ -432,7 +432,7 @@ async def acompletion(
messages: list = [],
functions: list | None = None,
function_call: str | None = None,
timeout: float | None = None,
timeout: float | httpx.Timeout | openai.Timeout | None = None,
temperature: float | None = None,
top_p: float | None = None,
n: int | None = None,
@ -5225,7 +5225,7 @@ def completion(
model: str,
# Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create
messages: list = [],
timeout: float | str | httpx.Timeout | None = None,
timeout: float | str | httpx.Timeout | openai.Timeout | None = None,
temperature: float | None = None,
top_p: float | None = None,
n: int | None = None,
@ -7880,6 +7880,8 @@ def adapter_completion(*, adapter_id: str, **kwargs) -> BaseModel | AdapterCompl
def moderation(input: str, model: str | None = None, api_key: str | None = None, **kwargs) -> OpenAIModerationResponse:
from litellm.llms.openai.common_utils import OpenAIHTTPClient
# only supports open ai for now
api_key = api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY")
@ -7888,10 +7890,7 @@ def moderation(input: str, model: str | None = None, api_key: str | None = None,
openai_client = kwargs.get("client", None)
if openai_client is None:
if api_base is not None:
openai_client = openai.OpenAI(api_key=api_key, base_url=api_base)
else:
openai_client = openai.OpenAI(api_key=api_key)
openai_client = openai.OpenAI(api_key=api_key, base_url=api_base, http_client=OpenAIHTTPClient())
if model is not None:
response = openai_client.moderations.create(input=input, model=model)
@ -8364,7 +8363,7 @@ def speech(
project: str | None = None,
max_retries: int | None = None,
metadata: dict | None = None,
timeout: float | httpx.Timeout | None = None,
timeout: float | httpx.Timeout | openai.Timeout | None = None,
response_format: str | None = None,
speed: int | None = None,
instructions: str | None = None,
@ -8394,9 +8393,12 @@ def speech(
if instructions is not None:
optional_params["instructions"] = instructions
if timeout is None:
timeout = litellm.request_timeout
timeout_or_default: Final = litellm.request_timeout if timeout is None else timeout
http_timeout: Final = (
CompletionTimeout.normalize(timeout_or_default)
if isinstance(timeout_or_default, openai.Timeout)
else timeout_or_default
)
if max_retries is None:
max_retries = litellm.num_retries or openai.DEFAULT_MAX_RETRIES
litellm_params_dict: Final = get_litellm_params(metadata=metadata, api_key=api_key or dynamic_api_key, **kwargs)
@ -8445,7 +8447,7 @@ def speech(
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
timeout=http_timeout,
extra_headers=extra_headers,
client=client,
_is_async=aspeech or False,
@ -8501,7 +8503,7 @@ def speech(
organization=organization,
project=project,
max_retries=max_retries,
timeout=timeout,
timeout=http_timeout,
logging_obj=logging_obj,
client=client, # pass AsyncOpenAI, OpenAI client
aspeech=aspeech,
@ -8532,7 +8534,7 @@ def speech(
optional_params=optional_params,
litellm_params_dict=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
timeout=http_timeout,
extra_headers=extra_headers,
base_llm_http_handler=base_llm_http_handler,
aspeech=aspeech or False,
@ -8580,7 +8582,7 @@ def speech(
azure_ad_token_provider=azure_ad_token_provider,
organization=organization,
max_retries=max_retries,
timeout=timeout,
timeout=http_timeout,
logging_obj=logging_obj,
client=client, # pass AsyncOpenAI, OpenAI client
aspeech=aspeech,
@ -8625,7 +8627,7 @@ def speech(
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
timeout=http_timeout,
extra_headers=extra_headers,
client=client,
_is_async=aspeech or False,
@ -8682,7 +8684,7 @@ def speech(
optional_params=optional_params,
litellm_params_dict=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
timeout=http_timeout,
extra_headers=headers,
base_llm_http_handler=base_llm_http_handler,
aspeech=aspeech or False,
@ -8728,7 +8730,7 @@ def speech(
optional_params=optional_params,
litellm_params_dict=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
timeout=http_timeout,
extra_headers=extra_headers,
base_llm_http_handler=base_llm_http_handler,
aspeech=aspeech or False,
@ -8769,7 +8771,7 @@ def speech(
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
timeout=http_timeout,
extra_headers=extra_headers,
client=client,
_is_async=aspeech or False,
@ -8797,7 +8799,7 @@ def speech(
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
timeout=http_timeout,
extra_headers=extra_headers,
client=client,
_is_async=aspeech or False,
@ -8821,7 +8823,7 @@ def speech(
optional_params=optional_params,
litellm_params_dict=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
timeout=http_timeout,
extra_headers=extra_headers,
base_llm_http_handler=base_llm_http_handler,
aspeech=aspeech or False,

2
litellm/model_prices_and_context_window_backup.json
View file

@ -71768,7 +71768,7 @@
"input_cost_per_token": 4.2e-08,
"litellm_provider": "azure_ai",
"mode": "evaluation",
"output_cost_per_token": 4.2e-08,
"output_cost_per_token": 0.0,
"source": "https://prices.azure.com/api/retail/prices?$filter=productName eq 'Microsoft Decision Models'"
},
"gpt-rosalind-discovery": {

7
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -8,7 +8,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, cast
from fastapi import HTTPException
from pydantic import TypeAdapter
from pydantic import SecretStr, TypeAdapter
from starlette.datastructures import Headers
from starlette.requests import Request
from starlette.types import Scope
@ -1532,6 +1532,11 @@ class MCPRequestHandler:
_get_mcp_server_auth_headers_from_headers = get_mcp_server_auth_headers_from_headers
@staticmethod
def get_incoming_bearer_token(headers: Mapping[str, str]) -> SecretStr | None:
authorization: Final = next((value for key, value in headers.items() if key.lower() == "authorization"), "")
return SecretStr(authorization[len("bearer ") :]) if authorization.lower().startswith("bearer ") else None
@staticmethod
def get_oauth2_headers_from_headers(headers: Headers) -> dict[str, str]:
"""

4
litellm/proxy/_experimental/mcp_server/contracts.py
View file

@ -4,10 +4,11 @@ from collections.abc import Mapping
from copy import deepcopy
from dataclasses import dataclass, field
from datetime import datetime
from types import MappingProxyType
from types import EllipsisType, MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias
from mcp.types import ErrorData, InputRequest, InputResponse, InputResponses
from pydantic import SecretStr
from litellm.proxy._experimental.mcp_server.tool_outcome import WireCompat
from litellm.proxy._types import UserAPIKeyAuth
@ -69,6 +70,7 @@ class OperationContext:
mcp_proxy_mode: bool = False
wire_compat: WireCompat = WireCompat.LEGACY
protocol_version: str | None = None
incoming_bearer_token: SecretStr | None | EllipsisType = field(default=..., repr=False)
def __post_init__(self) -> None:
object.__setattr__(self, "_caller", copy_caller(self._caller))

19
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -57,7 +57,7 @@ from mcp.types import (
ResourceTemplate,
)
from mcp.types import Tool as MCPTool
from pydantic import AnyUrl, Field, TypeAdapter
from pydantic import AnyUrl, Field, SecretStr, TypeAdapter
from typing_extensions import ReadOnly, assert_never
import litellm
@ -5921,6 +5921,7 @@ class MCPServerManager:
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
guardrail_context: Mapping[str, object] | None = None,
tool: MCPTool | None = None,
incoming_bearer_token: SecretStr | None | EllipsisType = ...,
) -> dict[str, Any]:
"""
Run pre-call checks and guardrail hooks for an MCP tool call.
@ -5977,13 +5978,11 @@ class MCPServerManager:
if proxy_logging_obj is None:
return hook_result
# Extract incoming Bearer token from raw request headers so
# guardrails like MCPJWTSigner can verify + re-sign it (FR-5).
normalized_raw: Final = {k.lower(): v for k, v in (raw_headers or {}).items()}
incoming_bearer_token: str | None = None
auth_hdr: Final = normalized_raw.get("authorization", "")
if auth_hdr.lower().startswith("bearer "):
incoming_bearer_token = auth_hdr[len("bearer ") :]
caller_bearer: Final = (
MCPRequestHandler.get_incoming_bearer_token(raw_headers or {})
if incoming_bearer_token is ...
else incoming_bearer_token
)
pre_hook_kwargs: Final = {
"guardrail_context": guardrail_context,
@ -5998,7 +5997,7 @@ class MCPServerManager:
getattr(user_api_key_auth, "end_user_id", None) if user_api_key_auth else None
),
"user_api_key_hash": (getattr(user_api_key_auth, "api_key_hash", None) if user_api_key_auth else None),
"incoming_bearer_token": incoming_bearer_token,
"incoming_bearer_token": caller_bearer.get_secret_value() if caller_bearer is not None else None,
"headers": logging_safe_mcp_headers(raw_headers),
"tool_description": tool.description if tool is not None else None,
"tool_input_schema": tool.input_schema if tool is not None else None,
@ -6664,6 +6663,7 @@ class MCPServerManager:
*,
catalog_auth_header: str | None | EllipsisType = ...,
listed_tool: MCPTool | None | EllipsisType = ...,
incoming_bearer_token: SecretStr | None | EllipsisType = ...,
) -> CallToolResult | InputRequiredResult:
"""
Call a tool with the given name and arguments
@ -6718,6 +6718,7 @@ class MCPServerManager:
litellm_logging_obj=litellm_logging_obj,
guardrail_context=guardrail_context,
tool=self.get_listed_tool(mcp_server, name, listed_caller) if listed_tool is ... else listed_tool,
incoming_bearer_token=incoming_bearer_token,
)
if "arguments" in hook_result:
arguments = hook_result["arguments"]

21
litellm/proxy/_experimental/mcp_server/operations.py
View file

@ -10,6 +10,7 @@ from contextvars import ContextVar
from dataclasses import dataclass, replace
from datetime import datetime
from functools import partial
from types import EllipsisType
from typing import Any, Final, NoReturn, TypeAlias, overload
from fastapi import HTTPException
@ -45,7 +46,7 @@ from mcp.types import (
ToolsCapability,
)
from mcp.types import Tool as MCPTool
from pydantic import AnyUrl, ConfigDict, Field, TypeAdapter
from pydantic import AnyUrl, ConfigDict, Field, SecretStr, TypeAdapter
from typing_extensions import ReadOnly, TypedDict, assert_never
from litellm._logging import verbose_logger
@ -352,6 +353,7 @@ async def _dispatch_virtual_mcp_tool(
oauth2_headers: dict[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
mcp_proxy_mode: bool = False,
incoming_bearer_token: SecretStr | None | EllipsisType = ...,
) -> CallToolResult | None:
"""Handle the mcp_tool_search / mcp_tool_call virtual tools.
@ -407,6 +409,7 @@ async def _dispatch_virtual_mcp_tool(
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_logging_obj=proxy_logging_obj,
incoming_bearer_token=incoming_bearer_token,
)
except Exception as exc:
if proxy_logging_obj is not None:
@ -509,6 +512,7 @@ async def _dispatch_virtual_mcp_tool(
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_logging_obj=virtual_logging_obj,
incoming_bearer_token=incoming_bearer_token,
)
@ -1874,6 +1878,7 @@ async def execute_mcp_tool(
wire_compat: WireCompat = WireCompat.LEGACY,
input_responses: InputResponses | None = None,
request_state: str | None = None,
incoming_bearer_token: SecretStr | None | EllipsisType = ...,
**kwargs: object, # kwargs-ok: preserves the existing REST and decorated logging call contract
) -> CallToolResult | InputRequiredResult:
context: Final = prepare_context(
@ -1884,6 +1889,7 @@ async def execute_mcp_tool(
raw_headers=raw_headers,
client_ip=client_ip,
wire_compat=wire_compat,
incoming_bearer_token=incoming_bearer_token,
)
operation: Final = AuthorizedToolCall(
name=name,
@ -1915,6 +1921,7 @@ async def _execute_mcp_tool(
wire_compat: WireCompat = WireCompat.LEGACY,
input_responses: InputResponses | None = None,
request_state: str | None = None,
incoming_bearer_token: SecretStr | None | EllipsisType = ...,
**kwargs: Any,
) -> CallToolResult | InputRequiredResult:
"""
@ -2152,6 +2159,7 @@ async def _execute_mcp_tool(
oauth2_headers,
),
),
incoming_bearer_token=incoming_bearer_token,
)
# `pre_call_tool_check` may return guardrail-modified
# arguments; honor them on the local path too.
@ -2212,6 +2220,7 @@ async def _execute_mcp_tool(
wire_compat=wire_compat,
input_responses=input_responses,
request_state=request_state,
incoming_bearer_token=incoming_bearer_token,
)
# Fall back to local tool registry with original name (legacy support)
@ -2284,6 +2293,7 @@ async def _execute_mcp_tool(
oauth2_headers,
),
),
incoming_bearer_token=incoming_bearer_token,
)
if "arguments" in hook_result:
arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args
@ -2455,6 +2465,7 @@ async def call_mcp_tool(
wire_compat: WireCompat = WireCompat.LEGACY,
input_responses: InputResponses | None = None,
request_state: str | None = None,
incoming_bearer_token: SecretStr | None | EllipsisType = ...,
**kwargs: Any,
) -> CallToolResult | InputRequiredResult:
"""
@ -2515,6 +2526,7 @@ async def call_mcp_tool(
wire_compat=wire_compat,
input_responses=input_responses,
request_state=request_state,
incoming_bearer_token=incoming_bearer_token,
**kwargs,
)
except Exception as e:
@ -2709,6 +2721,7 @@ async def _handle_managed_mcp_tool(
request_state: str | None = None,
*,
catalog_auth_header: str | None,
incoming_bearer_token: SecretStr | None | EllipsisType = ...,
) -> CallToolResult | InputRequiredResult:
"""Handle tool execution for managed server tools. ``catalog_auth_header`` is the header the client
supplied, which keys the caller's catalog slot; ``mcp_auth_header`` may already be the resolved
@ -2734,6 +2747,7 @@ async def _handle_managed_mcp_tool(
wire_compat=wire_compat,
input_responses=input_responses,
request_state=request_state,
incoming_bearer_token=incoming_bearer_token,
)
verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result)
return call_tool_result
@ -2908,6 +2922,7 @@ async def _execute_mcp_server_tool_call(
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
mcp_proxy_mode=context.mcp_proxy_mode,
incoming_bearer_token=context.incoming_bearer_token,
)
if virtual_tool_result is not None:
return virtual_tool_result
@ -2957,6 +2972,7 @@ async def _execute_mcp_server_tool_call(
wire_compat=context.wire_compat,
input_responses=params.input_responses,
request_state=params.request_state,
incoming_bearer_token=context.incoming_bearer_token,
**data, # for logging
)
except MCPMissingUserEnvVarsError as e:
@ -3184,6 +3200,7 @@ def prepare_context(
mcp_proxy_mode: bool = False,
wire_compat: WireCompat = WireCompat.LEGACY,
protocol_version: str | None = None,
incoming_bearer_token: SecretStr | None | EllipsisType = ...,
) -> OperationContext:
return OperationContext(
_caller=user_api_key_auth,
@ -3196,6 +3213,7 @@ def prepare_context(
mcp_proxy_mode=mcp_proxy_mode,
wire_compat=wire_compat,
protocol_version=protocol_version,
incoming_bearer_token=incoming_bearer_token,
)
@ -3376,6 +3394,7 @@ class GatewayOperations:
wire_compat=context.wire_compat,
input_responses=operation.input_responses,
request_state=operation.request_state,
incoming_bearer_token=context.incoming_bearer_token,
**operation.logging_data,
)
case ListToolsRequest(params=params):

3
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -229,6 +229,7 @@ if MCP_AVAILABLE:
from litellm.llms.litellm_proxy.skills.skill_search import (
DEFAULT_SKILL_SEARCH_TOP_K,
)
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: F401 # legacy module exports
_UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES,
@ -411,6 +412,7 @@ if MCP_AVAILABLE:
raw_headers=virtual_raw_headers,
litellm_logging_obj=virtual_logging_obj,
guardrail_context=MCPRequestContext.resolve_guardrail_context(virtual_data),
incoming_bearer_token=MCPRequestHandler.get_incoming_bearer_token(request.headers),
)
except Exception as e:
virtual_request_data: Final = virtual_processor.data
@ -1344,6 +1346,7 @@ if MCP_AVAILABLE:
litellm_logging_obj=data.get("litellm_logging_obj"),
guardrail_context=MCPRequestContext.resolve_guardrail_context(data),
requested_server_id=canonical_server_id,
incoming_bearer_token=MCPRequestHandler.get_incoming_bearer_token(request.headers),
)
result: Final = complete_call_tool_result(executed, WireCompat.LEGACY)
except Exception as e:

6
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -838,6 +838,7 @@ if MCP_AVAILABLE:
headers,
client_ip,
) = await get_or_extract_auth_context()
request: Final[object] = getattr(ctx, "request", None)
yield operations.prepare_context(
auth,
token,
@ -849,6 +850,11 @@ if MCP_AVAILABLE:
_mcp_proxy_mode.get(),
wire_compat_for(ctx.protocol_version),
ctx.protocol_version,
incoming_bearer_token=(
MCPRequestHandler.get_incoming_bearer_token(request.headers)
if isinstance(request, StarletteRequest)
else None
),
)
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListToolsResult:

8
litellm/proxy/_experimental/mcp_server/tool_search.py
View file

@ -8,13 +8,13 @@ from collections.abc import Iterator, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from itertools import chain, islice
from types import MappingProxyType
from types import EllipsisType, MappingProxyType
from typing import TYPE_CHECKING, Any, Final, TypedDict
import anyio
from anyio import to_process
from anyio.lowlevel import RunVar
from pydantic import JsonValue, ValidationError
from pydantic import JsonValue, SecretStr, ValidationError
from typing_extensions import ReadOnly, Required, assert_never
import litellm
@ -588,6 +588,7 @@ async def handle_mcp_proxy_tool(
oauth2_headers: dict[str, str] | None = None, # mutable-ok: preserve forwarded headers
raw_headers: dict[str, str] | None = None, # mutable-ok: preserve request headers
litellm_logging_obj: LiteLLMLoggingObj | None = None,
incoming_bearer_token: SecretStr | None | EllipsisType = ...,
) -> CallToolResult:
from fastapi import HTTPException
@ -662,6 +663,7 @@ async def handle_mcp_proxy_tool(
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
incoming_bearer_token=incoming_bearer_token,
)
@ -678,6 +680,7 @@ async def handle_mcp_tool_call(
litellm_logging_obj: LiteLLMLoggingObj | None = None,
requested_server_id: str | None = None,
guardrail_context: Mapping[str, object] | None = None,
incoming_bearer_token: SecretStr | None | EllipsisType = ...,
) -> CallToolResult:
from litellm.proxy._experimental.mcp_server.operations import (
_get_allowed_mcp_servers,
@ -719,5 +722,6 @@ async def handle_mcp_tool_call(
litellm_logging_obj=litellm_logging_obj,
requested_server_id=requested_server_id,
guardrail_context=guardrail_context,
incoming_bearer_token=incoming_bearer_token,
)
return complete_call_tool_result(result, WireCompat.LEGACY)

8
litellm/proxy/_types.py
View file

@ -3130,6 +3130,14 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="If True, router fallbacks configured in router_settings are only attempted when the calling key (and its team and project) is allowed to call the fallback model; unauthorized fallback targets are skipped and the primary model's error is returned. Default is False.",
)
disable_fallbacks_on_per_model_rate_limits: bool | None = Field(
None,
description=(
"If true, a request rejected by a key/team/org/project per-model rate limit "
"(model_rpm_limit / model_tpm_limit) returns 429 instead of retrying on the "
"configured fallbacks"
),
)
scheduled_job_stagger: ScheduledJobStaggerSettings | None = Field(
None,
description=(

2
litellm/proxy/auth/auth_utils.py
View file

@ -412,6 +412,8 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
"anthropic_workspace_id",
"anthropic-workspace-id",
"bedrock_tags",
# Deployment-pinned AgentCore session: every new caller-chosen value provisions another microVM.
"agentcore_runtime_session_id",
# Provider-specific endpoint overrides that flow into the outbound
# request via ``optional_params``. Same threat as ``api_base``:
# ``s3_endpoint_url`` redirects Bedrock file uploads to attacker

22
litellm/proxy/common_request_processing.py
View file

@ -16,6 +16,7 @@ from typing import (
Protocol,
TypeAlias,
TypeVar,
cast, # noqa: TID251 # _pre_call_with_fallbacks receives general_settings as a legacy bare dict
overload,
runtime_checkable,
)
@ -2302,7 +2303,21 @@ class ProxyBaseLLMRequestProcessing:
route_type: str,
llm_router: Router | None,
) -> tuple[dict, LiteLLMLoggingObj]:
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.common_utils.proxy_rate_limit_error import (
PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS,
ProxyRateLimitError,
per_model_rate_limits_disable_fallbacks,
)
general_settings_view: Final = cast( # cast-ok: this method keeps its legacy bare-dict settings parameter
Mapping[str, object], general_settings
)
def is_hard_per_model_limit(exc: ProxyRateLimitError) -> bool:
return (
exc.descriptor_key in PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS
and per_model_rate_limits_disable_fallbacks(general_settings_view)
)
configured_fallbacks: Final = (
self._configured_fallbacks(llm_router=llm_router, user_api_key_dict=user_api_key_dict)
@ -2336,6 +2351,7 @@ class ProxyBaseLLMRequestProcessing:
or not configured_fallbacks
or rate_limited_data.get("disable_fallbacks")
or not isinstance(original_model, str)
or is_hard_per_model_limit(original_exc)
):
raise
@ -2375,7 +2391,9 @@ class ProxyBaseLLMRequestProcessing:
llm_router=llm_router,
rate_limited_model=original_model,
)
except ProxyRateLimitError:
except ProxyRateLimitError as fallback_exc:
if is_hard_per_model_limit(fallback_exc):
raise
continue
except BaseException:
self.data = rate_limited_data

31
litellm/proxy/common_utils/proxy_rate_limit_error.py
View file

@ -40,9 +40,37 @@ from collections.abc import Mapping
from typing import Any, Final
from fastapi import HTTPException
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import RateLimitError, RateLimitErrorCategory, RateLimitType
PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS: Final = frozenset(
{
"model_per_key",
"model_per_team",
"model_per_organization",
"model_per_project",
"model_per_project_itpm",
"model_per_project_otpm",
}
)
DISABLE_FALLBACKS_ON_PER_MODEL_RATE_LIMITS_SETTING: Final = "disable_fallbacks_on_per_model_rate_limits"
_DISABLE_FALLBACKS_ON_PER_MODEL_RATE_LIMITS_FLAG: Final[TypeAdapter[bool | None]] = TypeAdapter(bool | None)
def per_model_rate_limits_disable_fallbacks(general_settings: Mapping[str, object]) -> bool:
raw_value: Final = general_settings.get(DISABLE_FALLBACKS_ON_PER_MODEL_RATE_LIMITS_SETTING)
try:
return _DISABLE_FALLBACKS_ON_PER_MODEL_RATE_LIMITS_FLAG.validate_python(raw_value) is True
except ValidationError:
verbose_proxy_logger.warning(
"general_settings.%s=%r is not a boolean, treating it as disabled",
DISABLE_FALLBACKS_ON_PER_MODEL_RATE_LIMITS_SETTING,
raw_value,
)
return False
def map_v3_rate_limit_type(
v3_value: str | None,
@ -149,6 +177,8 @@ class ProxyRateLimitError(HTTPException, RateLimitError):
rate_limit_type: str | RateLimitType | None = None,
model: str | None = None,
llm_provider: str | None = "litellm_proxy",
*,
descriptor_key: str | None = None,
):
# Normalize None → safe defaults so callers (and the resolver helper
# in `rate_limiter_utils`) can pass `None` without producing an
@ -191,3 +221,4 @@ class ProxyRateLimitError(HTTPException, RateLimitError):
self.headers = stringified_headers
self.detail = detail
self.status_code = 429
self.descriptor_key = descriptor_key

7
litellm/proxy/decisions_endpoints/endpoints.py
View file

@ -1,4 +1,5 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import Annotated, Final
from fastapi import APIRouter, Depends, Request, Response
@ -38,6 +39,10 @@ async def _invalid_request(
)
def _fields_checked_before_routing(data: Mapping[str, object]) -> Mapping[str, object]:
return MappingProxyType({key: value for key, value in data.items() if key != "safety_identifier"})
async def _request_data(request: Request, user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]:
body: Final = await request.body()
try:
@ -75,7 +80,7 @@ async def _process_decisions(
data: Final = await _request_data(request, user_api_key_dict)
try:
body_adapter.validate_python(data)
body_adapter.validate_python(_fields_checked_before_routing(data))
except ValidationError as error:
raise await _invalid_request(raw_data=data, error=error, user_api_key_dict=user_api_key_dict)
general_settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(proxy_general_settings)

11
litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
View file

@ -42,9 +42,10 @@ from litellm.proxy.spend_tracking.compression_savings import HEADROOM_GUARDRAIL_
from litellm.secret_managers.main import get_secret_str
from litellm.types.guardrails import GuardrailEventHooks, Mode
from litellm.types.integrations.custom_logger import (
HEADROOM_CONVERTED_STREAM_KEY,
HEADROOM_INTERCEPTION_PREFIX,
AgenticLoopPlan,
AgenticLoopRequestPatch,
as_converted_stream,
)
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
@ -883,18 +884,14 @@ class HeadroomGuardrail(CustomGuardrail):
call_type: CallTypes | None,
) -> dict[str, object] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict
base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type)
effective: Final = base_result if base_result is not None else kwargs
if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES:
return base_result
effective: Final = kwargs if base_result is None else _REQUEST_DATA_ADAPTER.validate_python(base_result)
if not effective.get("stream") or effective.get("background"):
return base_result
if not has_headroom_retrieve_tool(effective.get("tools")):
return base_result
return {
**effective,
"stream": False,
HEADROOM_CONVERTED_STREAM_KEY: True,
}
return as_converted_stream(effective, HEADROOM_INTERCEPTION_PREFIX)
async def async_should_run_agentic_loop(
self,

5
litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py
View file

@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import os
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, Literal, Protocol
@ -402,7 +403,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
verbose_proxy_logger.debug(
"Unable to authenticate to Hiddenlayer, JWT token is invalid or expired, trying to refresh the token."
)
self.jwt_token = self.refresh_jwt_func()
self.jwt_token = await asyncio.to_thread(self.refresh_jwt_func)
headers["Authorization"] = f"Bearer {self.jwt_token}"
response = await self._http_client.post(
f"{self.api_base}/detection/v1/interactions",
@ -609,7 +610,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
verbose_proxy_logger.debug(
"Unable to authenticate to Hiddenlayer, JWT token is invalid or expired, trying to refresh the token."
)
self.jwt_token = self.refresh_jwt_func()
self.jwt_token = await asyncio.to_thread(self.refresh_jwt_func)
headers["Authorization"] = f"Bearer {self.jwt_token}"
response = await self._http_client.post(
f"{self.api_base}/{path}",

4
litellm/proxy/hooks/parallel_request_limiter.py
View file

@ -105,6 +105,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current limits: max_parallel_requests: {max_parallel_requests}, tpm_limit: {tpm_limit}, rpm_limit: {rpm_limit}",
rate_limit_type=triggered_type,
requested_model=data.get("model") if data else None,
descriptor_key=rate_limit_type,
)
new_val = {
"current_requests": 1,
@ -143,6 +144,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
rate_limit_type=triggered_type,
model=resolved_model,
llm_provider=llm_provider,
descriptor_key=rate_limit_type,
)
await self.internal_usage_cache.async_batch_set_cache(
@ -170,6 +172,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
additional_details: str | None = None,
rate_limit_type: RateLimitType | None = None,
requested_model: str | None = None,
descriptor_key: str | None = None,
) -> NoReturn:
"""
Raise a 429 with a retry-after header for litellm-proxy parallel-request limits.
@ -207,6 +210,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
rate_limit_type=rate_limit_type or RateLimitType.CONCURRENT_REQUESTS,
model=resolved_model,
llm_provider=llm_provider,
descriptor_key=descriptor_key,
)
@with_service_target("rate_limits")

1
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -3594,6 +3594,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
rate_limit_type=map_v3_rate_limit_type(status["rate_limit_type"]),
model=resolved_model,
llm_provider=llm_provider,
descriptor_key=descriptor_key,
)
@staticmethod

10
litellm/proxy/litellm_pre_call_utils.py
View file

@ -223,6 +223,7 @@ def sanitize_for_log(value: object) -> str:
from litellm.router import Router
from litellm.secret_managers.main import get_secret_bool
from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES
from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
from litellm.types.services import ServiceTypes
from litellm.types.utils import (
@ -342,14 +343,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = (
# re-wrapped as a synthetic stream the caller never asked for, or raise the
# loop ceiling to drive many upstream model calls and sandbox executions
# from a single request.
"_agentic_loop_depth",
"_agentic_loop_fingerprints",
"_code_interpreter_interception_active",
"_code_interpreter_interception_converted_stream",
"_code_interpreter_interception_sandbox_key",
"_code_interpreter_interception_session_scoped",
"_headroom_interception_converted_stream",
"max_agentic_loops",
*AGENTIC_LOOP_KWARG_NAMES,
# Recomputed below from the actual caller-controlled timeout sources (headers and
# body fields); a client-forged value here would let a request either dodge cooldown
# protection on a real deployment failure or force a false "not caller-controlled"

3
litellm/proxy/proxy_server.py
View file

@ -12919,7 +12919,7 @@ async def audio_speech(
requested_format: Final = data.get("response_format")
upstream_content_type: Final = (
response.response.headers.get("content-type") if isinstance(response, HttpxBinaryResponseContent) else None
response.content_type if isinstance(response, HttpxBinaryResponseContent) else None
)
media_type: Final = resolve_speech_media_type(
upstream_content_type=upstream_content_type,
@ -18399,6 +18399,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
"pass_through_endpoints": "PydanticModel",
"store_model_in_db": "Boolean",
"store_prompts_in_spend_logs": "Boolean",
"disable_fallbacks_on_per_model_rate_limits": "Boolean",
"maximum_spend_logs_retention_period": "String",
"maximum_health_check_retention_period": "String",
"maximum_daily_tag_spend_retention_period": "String",

29
litellm/types/integrations/custom_logger.py
View file

@ -10,10 +10,20 @@ CHAT_COMPLETION_AGENTIC_SURFACE: Final[AgenticSurface] = "chat_completions"
RESPONSES_AGENTIC_SURFACE: Final[AgenticSurface] = "responses"
CODE_INTERPRETER_INTERCEPTION_PREFIX: Final = "_code_interpreter_interception"
HEADROOM_INTERCEPTION_PREFIX: Final = "_headroom_interception"
WEBSEARCH_INTERCEPTION_PREFIX: Final = "_websearch_interception"
HEADROOM_CONVERTED_STREAM_KEY: Final = f"{HEADROOM_INTERCEPTION_PREFIX}_converted_stream"
WEBSEARCH_CONVERTED_STREAM_KEY: Final = f"{WEBSEARCH_INTERCEPTION_PREFIX}_converted_stream"
WEBSEARCH_STREAM_OPTIONS_KEY: Final = f"{WEBSEARCH_INTERCEPTION_PREFIX}_stream_options"
CODE_INTERPRETER_STREAM_OPTIONS_KEY: Final = f"{CODE_INTERPRETER_INTERCEPTION_PREFIX}_stream_options"
HEADROOM_STREAM_OPTIONS_KEY: Final = f"{HEADROOM_INTERCEPTION_PREFIX}_stream_options"
STREAM_OPTIONS_STASH_KEYS: Final = (
CODE_INTERPRETER_STREAM_OPTIONS_KEY,
HEADROOM_STREAM_OPTIONS_KEY,
WEBSEARCH_STREAM_OPTIONS_KEY,
)
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES: Final = frozenset(
(
"_websearch_interception",
WEBSEARCH_INTERCEPTION_PREFIX,
"_compression_interception",
HEADROOM_INTERCEPTION_PREFIX,
)
@ -40,6 +50,23 @@ def converted_stream_requested(params: Mapping[str, object]) -> bool:
return any(bool(params.get(key)) for key in CONVERTED_STREAM_KEYS)
def as_converted_stream(
params: Mapping[str, object], prefix: str
) -> dict[str, object]: # mutable-ok: deployment hooks hand back the kwargs dict their caller keeps editing
stream_options: Final = params.get("stream_options")
stash: Final = {} if stream_options is None else {f"{prefix}_stream_options": stream_options}
return {
**{key: value for key, value in params.items() if key != "stream_options"},
"stream": False,
f"{prefix}_converted_stream": True,
**stash,
}
def stashed_stream_options(params: Mapping[str, object]) -> object:
return next((params[key] for key in STREAM_OPTIONS_STASH_KEYS if params.get(key) is not None), None)
class AgenticLoopSafetyError(ValueError):
"""
Raised when an agentic-loop safety rail refuses a rerun.

9
litellm/types/litellm_params.py
View file

@ -355,15 +355,24 @@ class AgenticLoopState:
code_interpreter_converted_stream: bool | None = field(
default=None, metadata=wire("_code_interpreter_interception_converted_stream")
)
code_interpreter_stream_options: Mapping[str, object] | None = field(
default=None, metadata=wire("_code_interpreter_interception_stream_options")
)
websearch_emit_native_blocks: bool | None = field(
default=None, metadata=wire("_websearch_interception_emit_native_blocks")
)
websearch_converted_stream: bool | None = field(
default=None, metadata=wire("_websearch_interception_converted_stream")
)
websearch_stream_options: Mapping[str, object] | None = field(
default=None, metadata=wire("_websearch_interception_stream_options")
)
headroom_converted_stream: bool | None = field(
default=None, metadata=wire("_headroom_interception_converted_stream")
)
headroom_stream_options: Mapping[str, object] | None = field(
default=None, metadata=wire("_headroom_interception_stream_options")
)
@dataclass(frozen=True, slots=True, kw_only=True)

26
litellm/types/llms/openai.py
View file

@ -2,13 +2,14 @@ import builtins
from collections.abc import Iterable, Mapping
from enum import Enum
from os import PathLike
from typing import IO, Any, Final, Literal, Optional, TypeAlias, Union
from typing import IO, Any, Final, Generic, Literal, Optional, TypeAlias, Union
import httpx
from openai import Omit
from openai._legacy_response import (
HttpxBinaryResponseContent as _HttpxBinaryResponseContent,
)
from openai._types import Response as SDKResponse
from openai.lib.streaming._assistants import (
AssistantEventHandler,
AssistantStreamManager,
@ -79,6 +80,7 @@ from typing_extensions import (
ReadOnly,
Required,
TypedDict,
TypeVar,
override,
)
@ -118,8 +120,12 @@ class BinaryResponseSummary(TypedDict):
num_bytes: ReadOnly[int]
class HttpxBinaryResponseContent(_HttpxBinaryResponseContent):
_ResponseT = TypeVar("_ResponseT", bound=httpx.Response | SDKResponse, default=httpx.Response)
class HttpxBinaryResponseContent(_HttpxBinaryResponseContent, Generic[_ResponseT]):
_hidden_params: dict
response: _ResponseT # pyright: ignore[reportIncompatibleVariableOverride] # SDK accepts both backends at runtime
@property
def hidden_params(self) -> dict[str, builtins.object]: # mutable-ok: API requires mutation
@ -129,22 +135,28 @@ class HttpxBinaryResponseContent(_HttpxBinaryResponseContent):
def hidden_params(self, hidden_params: dict[str, builtins.object]) -> None: # mutable-ok: API requires mutation
self._hidden_params = hidden_params
def __init__(self, response: httpx.Response) -> None:
super().__init__(response)
def __init__(self, response: _ResponseT) -> None:
super().__init__(response) # pyright: ignore[reportArgumentType] # SDK accepts both backends at runtime
self._hidden_params = {}
@property
def content_type(self) -> str | None:
headers: Final[Mapping[str, str]] = self.response.headers
return headers.get("content-type")
def logging_summary(self) -> BinaryResponseSummary:
return {
"object": "binary",
"content_type": self.response.headers.get("content-type"),
"content_type": self.content_type,
"num_bytes": self._num_bytes(),
}
def _num_bytes(self) -> int:
try:
return len(self.response.content)
except httpx.ResponseNotRead:
content: Final = self.response.content
except RuntimeError:
return self.response.num_bytes_downloaded
return len(content)
def set_response_cost(self, response_cost: float | None) -> None:
if response_cost is None:

13
litellm/types/router.py
View file

@ -21,11 +21,13 @@ from typing import (
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
import httpx
from openai import Timeout as SDKTimeout
from pydantic import ConfigDict, Field, JsonValue, field_validator, model_validator
from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_checkable
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.provider_affinity import validate_provider_affinity_header_name
from litellm.types.llms.base import LiteLLMBaseModel
@ -481,7 +483,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
rpm: int | None = None
itpm: int | None = None
otpm: int | None = None
timeout: float | str | httpx.Timeout | None = None # if str, pass in as os.environ/
timeout: float | str | httpx.Timeout | SDKTimeout | None = None # if str, pass in as os.environ/
stream_timeout: float | str | None = None # timeout when making stream=True calls, if str, pass in as os.environ/
max_retries: int | None = None
drop_params: bool | str | None = None
@ -564,6 +566,13 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
valkey_text_field: str | None = None
valkey_embedding_field: str | None = None
@field_validator("timeout")
@classmethod
def normalize_timeout(
cls, value: float | str | httpx.Timeout | SDKTimeout | None
) -> float | str | httpx.Timeout | None:
return CompletionTimeout.normalize(value) if isinstance(value, (httpx.Timeout, SDKTimeout)) else value
@field_validator("provider_affinity_header")
@classmethod
def validate_provider_affinity_header(cls, value: str | None) -> str | None:
@ -677,7 +686,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False):
api_key: str | None
api_base: str | None
api_version: str | None
timeout: float | str | httpx.Timeout | None
timeout: float | str | httpx.Timeout | SDKTimeout | None # writable-ok: preserve timeout updates
stream_timeout: float | str | None
max_retries: int | None
organization: list | str | None # for openai orgs

2
model_prices_and_context_window.json
View file

@ -71768,7 +71768,7 @@
"input_cost_per_token": 4.2e-08,
"litellm_provider": "azure_ai",
"mode": "evaluation",
"output_cost_per_token": 4.2e-08,
"output_cost_per_token": 0.0,
"source": "https://prices.azure.com/api/retail/prices?$filter=productName eq 'Microsoft Decision Models'"
},
"gpt-rosalind-discovery": {

2
packaging/litellm-core/pyproject.toml
View file

@ -13,7 +13,7 @@ dependencies = [
"fastuuid>=0.14.0,<1.0",
"filelock>=3.16.1,<4.0",
"httpx[http2]>=0.28.0,<1.0",
"openai>=2.20.0,<3.0.0",
"openai>=2.20.0,<4.0.0",
"python-dateutil>=2.8.2,<3.0",
"python-dotenv>=1.0.0,<2.0",
"pyyaml>=6.0.3,<7.0",

2
pyproject.toml
View file

@ -17,7 +17,7 @@ dependencies = [
"fastuuid>=0.14.0,<1.0",
"filelock>=3.16.1,<4.0",
"httpx[http2]>=0.28.0,<1.0",
"openai>=2.20.0,<3.0.0",
"openai>=2.20.0,<4.0.0",
"python-dotenv>=1.0.0,<2.0",
"pyyaml>=6.0.3,<7.0",
"packaging>=24.0",

41
tests/base_sdk_tests/check_base_sdk_install.py
View file

@ -8,11 +8,18 @@ the very class of undeclared-dependency bug this guards against.
import argparse
import importlib.util
import re
import sys
import traceback
from collections.abc import Callable
from importlib.metadata import PackageNotFoundError, packages_distributions, requires
from itertools import chain
from typing import Final
EXTRAS_ONLY_MODULES = ("fastapi", "uvicorn", "keyring", "mcp", "mcp_types", "httpx2", "httpcore2")
SDK_DISTRIBUTIONS: Final = frozenset({"litellm", "litellm-core"})
REQUIREMENT_NAME: Final = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*")
EXTRA_MARKER: Final = re.compile(r";.*\bextra\s*==")
def _require(condition: bool, message: str) -> None:
@ -20,13 +27,43 @@ def _require(condition: bool, message: str) -> None:
raise AssertionError(message)
def _canonical(name: str) -> str:
return re.sub(r"[-_.]+", "-", name).lower()
def _base_requirements(distribution: str) -> frozenset[str]:
try:
declared: Final = requires(distribution) or ()
except PackageNotFoundError:
return frozenset()
return frozenset(
_canonical(match.group(0))
for requirement in declared
if not EXTRA_MARKER.search(requirement) and (match := REQUIREMENT_NAME.match(requirement))
)
def _base_closure(pending: frozenset[str], reached: frozenset[str] = frozenset()) -> frozenset[str]:
if not pending:
return reached
expanded: Final = reached | pending
return _base_closure(frozenset(chain.from_iterable(_base_requirements(name) for name in pending)) - expanded, expanded)
def check_environment_is_base_only() -> str:
present = tuple(name for name in EXTRAS_ONLY_MODULES if importlib.util.find_spec(name) is not None)
from_base: Final = _base_closure(SDK_DISTRIBUTIONS)
owners: Final = packages_distributions()
installed: Final = tuple(name for name in EXTRAS_ONLY_MODULES if importlib.util.find_spec(name) is not None)
inherited: Final = tuple(
name for name in installed if from_base & {_canonical(owner) for owner in owners.get(name) or (name,)}
)
present: Final = tuple(name for name in installed if name not in inherited)
_require(
not present,
f"{', '.join(present)} installed, so this environment is not base-only and the run proves nothing",
)
return f"no extras-only packages present ({', '.join(EXTRAS_ONLY_MODULES)})"
from_dependencies: Final = f"; {', '.join(inherited)} come from base dependencies" if inherited else ""
return f"no extras-only packages present ({', '.join(EXTRAS_ONLY_MODULES)}){from_dependencies}"
def check_import() -> str:

7
tests/base_sdk_tests/check_sdk_http.py
View file

@ -37,7 +37,8 @@ class RecordingUpstream(BaseHTTPRequestHandler):
pass
def do_POST(self) -> None:
REQUESTS.put((self.path, dict(self.headers), self.rfile.read(int(self.headers["Content-Length"]))))
headers: Final = {name.lower(): value for name, value in self.headers.items()}
REQUESTS.put((self.path, headers, self.rfile.read(int(self.headers["Content-Length"]))))
status, body = RESPONSES.get(timeout=10)
ARRIVED.set()
if status == 0:
@ -78,7 +79,7 @@ def check_http(base: str) -> None:
assert response.usage.total_tokens == 5
path, headers, body = REQUESTS.get(timeout=10)
assert path == "/v1/chat/completions"
assert headers["Authorization"] == "Bearer test-key"
assert headers["authorization"] == "Bearer test-key"
assert json.loads(body)["messages"] == MESSAGES
enqueue_stream()
assert "".join(part.choices[0].delta.content or "" for part in litellm.completion(**arguments, stream=True)) == "pong"
@ -156,7 +157,7 @@ def check_http(base: str) -> None:
assert result.usage.total_tokens == 5
path, headers, body = REQUESTS.get(timeout=10)
assert path.endswith("/converse")
assert {key.lower(): value for key, value in headers.items()}["authorization"].startswith(
assert headers["authorization"].startswith(
"AWS4-HMAC-SHA256" if signed else "Bearer bearer-key"
)
assert json.loads(body)["messages"][0]["content"][0]["text"] == "ping"

46
tests/integration/_support/forward_proxy.py
View file

@ -1,5 +1,7 @@
from __future__ import annotations
import contextlib
import socket
import threading
from collections.abc import Generator, Mapping
from contextlib import contextmanager
@ -69,3 +71,47 @@ def refusing_forward_proxy(port: int = 0) -> Generator[ForwardProxy, None, None]
thread.join(timeout=6)
assert not thread.is_alive(), "Owned forward proxy survived cleanup"
server.server_close()
def _pipe(source: socket.socket, sink: socket.socket) -> None:
with contextlib.suppress(OSError):
for chunk in iter(lambda: source.recv(65536), b""):
sink.sendall(chunk)
with contextlib.suppress(OSError):
sink.shutdown(socket.SHUT_WR)
@contextmanager
def tunnelling_forward_proxy(destination_port: int, port: int = 0) -> Generator[ForwardProxy, None, None]:
received: Final[SimpleQueue[Tunnel]] = SimpleQueue()
class Handler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
rbufsize = 0
timeout = 30
request: socket.socket
def do_CONNECT(self) -> None:
received.put(Tunnel(self.command, self.path, {name.lower(): value for name, value in self.headers.items()}))
self.close_connection = True
with socket.create_connection(("127.0.0.1", destination_port), timeout=30) as destination:
self.send_response(200, "Connection established")
self.end_headers()
outbound: Final = threading.Thread(target=_pipe, args=(self.request, destination))
outbound.start()
_pipe(destination, self.request)
outbound.join(timeout=35)
def log_message(self, format: str, *args: object) -> None:
pass
with ThreadingHTTPServer(("127.0.0.1", port), Handler) as server:
thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
thread.start()
try:
yield ForwardProxy(f"http://127.0.0.1:{server.server_port}", server.server_port, received)
finally:
server.shutdown()
thread.join(timeout=6)
assert not thread.is_alive(), "Owned forward proxy survived cleanup"
server.server_close()

245
tests/integration/_support/interception_vendor.py Normal file
View file

@ -0,0 +1,245 @@
from __future__ import annotations
import json
import re
import uuid
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Final, Literal, TypeAlias
from urllib.parse import urlsplit
from integration._support.openai_wire import chat_stream, responses_stream
from integration._support.wire import Reply, Request
from pydantic import JsonValue, TypeAdapter
AZURE_KEY: Final = "synthetic-azure-key"
TAVILY_KEY: Final = "synthetic-tavily-key"
API_VERSION: Final = "2025-04-01-preview"
SEARCH_TOOL: Final = "litellm_web_search"
SEARCH_TARGET: Final = "/tavily/search"
STREAM_OPTIONS_REFUSAL: Final = "The 'stream_options' parameter is only allowed when 'stream' is enabled."
FOLLOWUP_FAILURE: Final = "scripted follow-up failure"
FIRST_USAGE: Final = {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18}
FOLLOWUP_USAGE: Final = {"prompt_tokens": 23, "completion_tokens": 5, "total_tokens": 28}
RESPONSES_FIRST_USAGE: Final = {"input_tokens": 13, "output_tokens": 6, "total_tokens": 19}
RESPONSES_FOLLOWUP_USAGE: Final = {"input_tokens": 29, "output_tokens": 4, "total_tokens": 33}
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_ITEMS: Final = TypeAdapter(list[dict[str, JsonValue]])
_CHAT_TARGET: Final = re.compile(r"^/openai/deployments/([^/]+)/chat/completions$")
Mode: TypeAlias = Literal["search", "plain", "search-down", "followup-fails", "missing"]
def deployment(mode: Mode) -> str:
return f"dep-{uuid.uuid4().hex}-{mode}"
def answer(name: str) -> str:
return f"answer-{name}"
def plain(name: str) -> str:
return f"plain-{name}"
def query(name: str) -> str:
return f"q-{name}"
def _minted(prefix: str) -> str:
return f"{prefix}{uuid.uuid4().hex[:12]}"
def _azure_error(status: int, message: str, *, param: str | None = None) -> Reply:
return Reply(
status=status,
body=json.dumps(
{"error": {"message": message, "type": "invalid_request_error", "param": param, "code": None}}
).encode(),
)
def _refusal(body: Mapping[str, JsonValue]) -> Reply | None:
if body.get("stream_options") is not None and body.get("stream") is not True:
return _azure_error(400, STREAM_OPTIONS_REFUSAL, param="stream_options")
unknown: Final = tuple(key for key in body if key.startswith("_"))
if unknown:
return _azure_error(400, f"Unrecognized request argument supplied: {unknown[0]}")
return None
def _offers_search(body: Mapping[str, JsonValue]) -> bool:
tools: Final = _ITEMS.validate_python(body.get("tools") or [])
return any(
tool.get("name") == SEARCH_TOOL
or JSON_OBJECT.validate_python(tool.get("function") or {}).get("name") == SEARCH_TOOL
for tool in tools
)
def _chat_completion(identity: str, message: Mapping[str, JsonValue], finish: str, usage: Mapping[str, int]) -> Reply:
return Reply(
body=json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": "gpt-4o",
"choices": [{"index": 0, "message": dict(message), "finish_reason": finish}],
"usage": dict(usage),
}
).encode()
)
def _chat(name: str, body: Mapping[str, JsonValue]) -> Reply:
if body.get("stream") is True:
return Reply(
content_type="text/event-stream", chunks=chat_stream(_minted(f"chatcmpl-{name}-s-"), "gpt-4o", plain(name))
)
messages: Final = _ITEMS.validate_python(body["messages"])
if any(message.get("role") == "tool" for message in messages):
if name.endswith("-followup-fails"):
return _azure_error(500, FOLLOWUP_FAILURE)
return _chat_completion(
_minted(f"chatcmpl-{name}-2-"), {"role": "assistant", "content": answer(name)}, "stop", FOLLOWUP_USAGE
)
if name.endswith("-plain") or not _offers_search(body):
return _chat_completion(
_minted(f"chatcmpl-{name}-1-"), {"role": "assistant", "content": plain(name)}, "stop", FIRST_USAGE
)
call: Final = {
"id": f"call_{name[4:36]}",
"type": "function",
"function": {"name": SEARCH_TOOL, "arguments": json.dumps({"query": query(name)})},
}
message: Final = JSON_OBJECT.validate_python({"role": "assistant", "content": None, "tool_calls": [call]})
return _chat_completion(_minted(f"chatcmpl-{name}-1-"), message, "tool_calls", FIRST_USAGE)
def _response(identity: str, name: str, output: Sequence[Mapping[str, JsonValue]], usage: Mapping[str, int]) -> Reply:
return Reply(
body=json.dumps(
{
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed",
"model": name,
"output": [dict(item) for item in output],
"parallel_tool_calls": True,
"tool_choice": "auto",
"tools": [],
"usage": dict(usage),
}
).encode()
)
def _text_item(identity: str, text: str) -> Mapping[str, JsonValue]:
return {
"id": f"msg_{identity}",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": text, "annotations": []}],
}
def _responses(body: Mapping[str, JsonValue]) -> Reply:
name: Final = str(body["model"])
if body.get("stream") is True:
return Reply(
content_type="text/event-stream", chunks=responses_stream(_minted(f"resp_{name}_s_"), name, plain(name))
)
items: Final = _ITEMS.validate_python(body["input"]) if isinstance(body.get("input"), list) else []
if any(item.get("type") == "function_call_output" for item in items):
if name.endswith("-followup-fails"):
return _azure_error(500, FOLLOWUP_FAILURE)
return _response(
_minted(f"resp_{name}_2_"), name, (_text_item(f"{name}_2", answer(name)),), RESPONSES_FOLLOWUP_USAGE
)
if name.endswith("-plain") or not _offers_search(body):
return _response(
_minted(f"resp_{name}_1_"), name, (_text_item(f"{name}_1", plain(name)),), RESPONSES_FIRST_USAGE
)
call: Final = {
"type": "function_call",
"id": f"fc_{name[4:36]}",
"call_id": f"call_{name[4:36]}",
"name": SEARCH_TOOL,
"arguments": json.dumps({"query": query(name)}),
"status": "completed",
}
return _response(_minted(f"resp_{name}_1_"), name, (call,), RESPONSES_FIRST_USAGE)
def _search(request: Request, body: Mapping[str, JsonValue]) -> Reply:
if request.headers.get("authorization") != f"Bearer {TAVILY_KEY}":
return Reply(status=401, body=b'{"detail":{"error":"Unauthorized"}}')
text: Final = str(body.get("query", ""))
if text.endswith("-search-down"):
return Reply(status=500, body=b'{"detail":{"error":"scripted search outage"}}')
return Reply(
body=json.dumps(
{
"query": text,
"results": [
{"title": f"title-{text}", "url": "https://example.test/result", "content": f"snippet-{text}"}
],
}
).encode()
)
def respond(request: Request) -> Reply:
target: Final = urlsplit(request.target).path
if request.method == "GET":
return Reply(body=b'{"object":"list","data":[]}')
body: Final = JSON_OBJECT.validate_json(request.body or b"{}")
if target == SEARCH_TARGET:
return _search(request, body)
if request.headers.get("api-key") != AZURE_KEY:
return Reply(
status=401, body=b'{"error":{"code":"401","message":"Access denied due to invalid subscription key."}}'
)
chat: Final = _CHAT_TARGET.match(target)
if chat is not None and chat.group(1).endswith("-missing"):
return Reply(
status=404,
body=b'{"error":{"code":"DeploymentNotFound","message":"The API deployment for this resource does not exist."}}',
)
refused: Final = _refusal(body)
if refused is not None:
return refused
if chat is not None:
return _chat(chat.group(1), body)
if target.endswith("/responses"):
return _responses(body)
return Reply(status=404, body=b'{"error":{"code":"404","message":"Resource not found"}}')
@dataclass(frozen=True, slots=True)
class Received:
model_calls: tuple[Mapping[str, JsonValue], ...]
searches: tuple[Mapping[str, JsonValue], ...]
def _named(request: Request, name: str) -> bool:
return name in request.target or name in request.body.decode(errors="replace")
def received(drained: Sequence[Request], name: str) -> Received:
requests: Final = tuple(request for request in drained if request.method == "POST" and _named(request, name))
return Received(
model_calls=tuple(
JSON_OBJECT.validate_json(request.body)
for request in requests
if urlsplit(request.target).path != SEARCH_TARGET
),
searches=tuple(
JSON_OBJECT.validate_json(request.body)
for request in requests
if urlsplit(request.target).path == SEARCH_TARGET
),
)

1200
tests/integration/compatibility/test_openai_sdk_client_paths.py Normal file
View file

File diff suppressed because it is too large Load diff

424
tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py Normal file
View file

@ -0,0 +1,424 @@
from __future__ import annotations
import json
import signal
import subprocess
import uuid
from collections.abc import Mapping
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass
from functools import partial
from itertools import product
from pathlib import Path
from typing import Final, Literal, TypeAlias
import httpx
import psutil
import pytest
from anthropic import Anthropic
from anthropic import APIError as AnthropicApiError
from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value
from integration._support.database import read_rows
from integration._support.process import graceful_stop_seconds, owned_proxy_process, owned_upstream
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
from openai import APIError as OpenAiApiError
from openai import OpenAI
from pydantic import JsonValue
from tests.integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, SseResponse
_Kind: TypeAlias = Literal["chat", "chat_stream", "messages", "responses", "completions", "moderations"]
_ALL_KINDS: Final[tuple[_Kind, ...]] = ("chat", "chat_stream", "messages", "responses", "completions", "moderations")
_CHAT_KINDS: Final[tuple[_Kind, ...]] = ("chat", "chat_stream")
_NO_CACHE: Final = {"cache": {"no-cache": True}}
_SPEND_QUERY: Final = 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = %s'
_CHAT: Final = JsonResponse(
content_type="application/json",
body={
"id": "chatcmpl-$UNIQUE_ID",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
},
)
_MODERATION: Final = JsonResponse(
content_type="application/json",
body={
"id": "modr-$UNIQUE_ID",
"model": "omni-moderation-latest",
"results": [{"flagged": False, "categories": {}, "category_scores": {}}],
},
)
_RESPONSES: Final = JsonResponse(
content_type="application/json",
body={
"id": "resp_$UNIQUE_ID",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-4o-mini",
"output": [
{
"type": "message",
"id": "msg_$UNIQUE_ID",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "scripted", "annotations": []}],
}
],
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
},
)
_PLAIN: Final = RoutedResponse(
content_type="application/x-routed",
routes={"POST /chat/completions": _CHAT, "POST /moderations": _MODERATION, "POST /responses": _RESPONSES},
)
_STREAM: Final = SseResponse(
content_type="text/event-stream",
frames=(
'data: {"id":"chatcmpl-$UNIQUE_ID","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",'
'"choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}',
'data: {"id":"chatcmpl-$UNIQUE_ID","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",'
'"choices":[{"index":0,"delta":{"content":"streamed"},"finish_reason":null}]}',
'data: {"id":"chatcmpl-$UNIQUE_ID","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",'
'"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}',
"data: [DONE]",
),
frame_delay_ms=150,
)
@dataclass(frozen=True, slots=True)
class _Call:
kind: _Kind
marker: str
@dataclass(frozen=True, slots=True)
class _Completed:
call: _Call
response_id: str
@dataclass(frozen=True, slots=True)
class _Failed:
call: _Call
error: str
@dataclass(frozen=True, slots=True)
class _Clients:
openai: OpenAI
anthropic: Anthropic
plain_model: str
stream_model: str
class _Observations:
def __init__(self, url: str) -> None:
self.url = url.rstrip("/")
self.items: tuple[Mapping[str, JsonValue], ...] = ()
def read(self) -> tuple[Mapping[str, JsonValue], ...]:
with httpx.Client(timeout=10, trust_env=False) as client:
payload: Final = JSON_OBJECT.validate_python(client.get(f"{self.url}/__observations").json())
requests: Final = payload.get("requests")
assert isinstance(requests, list)
self.items = (*self.items, *(object_value(item) for item in requests if isinstance(item, dict)))
return self.items
def _register(
scenario: Scenario, upstream_url: str, label: str, response: RoutedResponse | SseResponse
) -> ScenarioHandle:
handle: Final = register_scenario(f"chaos-{label}-{uuid.uuid4().hex}", response, control_url=upstream_url)
scenario.cleanups.callback(delete_scenario, handle)
return handle
def _deployments(scenario: Scenario, upstream_url: str) -> tuple[str, str]:
plain_handle: Final = _register(scenario, upstream_url, "plain", _PLAIN)
stream_handle: Final = _register(scenario, upstream_url, "stream", _STREAM)
return (
scenario.model(api_base=plain_handle.api_base(), api_key=plain_handle.scenario_id),
scenario.model(api_base=stream_handle.api_base(), api_key=stream_handle.scenario_id),
)
def _clients(gateway: Gateway, deployments: tuple[str, str]) -> _Clients:
base: Final = str(gateway.client.base_url).rstrip("/")
clients: Final = _Clients(
openai=OpenAI(base_url=f"{base}/v1", api_key=gateway.key, max_retries=0),
anthropic=Anthropic(base_url=base, api_key=gateway.key, max_retries=0),
plain_model=deployments[0],
stream_model=deployments[1],
)
_import_lazy_sdk_resources_on_this_thread(clients)
return clients
def _import_lazy_sdk_resources_on_this_thread(clients: _Clients) -> None:
resources: Final = (
clients.openai.chat.completions,
clients.openai.completions,
clients.openai.responses,
clients.openai.moderations,
clients.anthropic.messages,
)
assert all(resource is not None for resource in resources)
def _invoke(clients: _Clients, call: _Call) -> str:
match call.kind:
case "chat":
return clients.openai.chat.completions.create(
model=clients.plain_model, messages=[{"role": "user", "content": call.marker}], extra_body=_NO_CACHE
).id
case "chat_stream":
chunks: Final = tuple(
clients.openai.chat.completions.create(
model=clients.stream_model,
messages=[{"role": "user", "content": call.marker}],
stream=True,
extra_body=_NO_CACHE,
)
)
assert chunks[-1].choices[0].finish_reason == "stop", chunks
assert len({chunk.id for chunk in chunks}) == 1, chunks
return chunks[0].id
case "messages":
return clients.anthropic.messages.create(
model=clients.plain_model,
max_tokens=64,
messages=[{"role": "user", "content": call.marker}],
extra_body=_NO_CACHE,
).id
case "responses":
return clients.openai.responses.create(
model=clients.plain_model, input=call.marker, extra_body=_NO_CACHE
).id
case "completions":
return clients.openai.completions.create(
model=clients.plain_model, prompt=call.marker, extra_body=_NO_CACHE
).id
case "moderations":
return clients.openai.moderations.create(
model=clients.plain_model, input=call.marker, extra_body=_NO_CACHE
).id
def _attempt(clients: _Clients, call: _Call) -> _Completed | _Failed:
try:
return _Completed(call, _invoke(clients, call))
except (OpenAiApiError, AnthropicApiError, httpx.HTTPError) as error:
return _Failed(call, f"{type(error).__name__}: {str(error)[:200]}")
def _every_worker_serves_every_kind(gateway: Gateway, deployments: tuple[str, str]) -> bool:
calls: Final = _burst_calls(_ALL_KINDS, 2)
with ThreadPoolExecutor(max_workers=len(calls)) as pool:
futures: Final = _submit(pool, _clients(gateway, deployments), calls)
outcomes: Final = tuple(future.result(timeout=60) for future in futures)
return all(isinstance(outcome, _Completed) for outcome in outcomes)
def _burst_calls(kinds: tuple[_Kind, ...], per_kind: int) -> tuple[_Call, ...]:
return tuple(_Call(kind, f"burst-{kind}-{uuid.uuid4().hex}") for kind, _ in product(kinds, range(per_kind)))
def _submit(
pool: ThreadPoolExecutor, clients: _Clients, calls: tuple[_Call, ...]
) -> tuple[Future[_Completed | _Failed], ...]:
return tuple(pool.submit(_attempt, clients, call) for call in calls)
def _done_count(futures: tuple[Future[_Completed | _Failed], ...]) -> int:
return sum(future.done() for future in futures)
def _liveliness(gateway: Gateway) -> int:
return gateway.client.get("/health/liveliness").status_code
def _spend_rows(request_id: str) -> tuple[Mapping[str, JsonValue], ...]:
return tuple(read_rows(_SPEND_QUERY, (request_id,)))
def _marker_counts(items: tuple[Mapping[str, JsonValue], ...], markers: tuple[str, ...]) -> tuple[int, ...]:
bodies: Final = tuple(json.dumps(item.get("body")) for item in items)
return tuple(sum(marker in body for body in bodies) for marker in markers)
def _landed_once(response_ids: tuple[str, ...]) -> tuple[tuple[Mapping[str, JsonValue], ...], ...]:
rows: Final = tuple(
eventually(partial(_spend_rows, response_id), lambda values: len(values) == 1, seconds=90)
for response_id in response_ids
)
assert all(found[0]["request_id"] == response_id for found, response_id in zip(rows, response_ids)), rows
return rows
def _is_worker(child: psutil.Process) -> bool:
try:
return "spawn_main" in " ".join(child.cmdline())
except psutil.Error:
return False
def _worker_alive(worker: psutil.Process) -> bool:
try:
return worker.is_running() and worker.status() != psutil.STATUS_ZOMBIE
except psutil.NoSuchProcess:
return False
def _alive_workers(process: subprocess.Popen[bytes]) -> tuple[psutil.Process, ...]:
children: Final = tuple(child for child in psutil.Process(process.pid).children() if _is_worker(child))
return tuple(child for child in children if _worker_alive(child))
def _process_tree(process: subprocess.Popen[bytes]) -> str:
root: Final = psutil.Process(process.pid)
return "\n".join(_process_line(member) for member in (root, *root.children(recursive=True)))
def _process_line(member: psutil.Process) -> str:
try:
return f"{member.pid} {' '.join(member.cmdline())}"
except psutil.Error:
return f"{member.pid} <exited>"
def _kind_counts(outcomes: tuple[_Completed | _Failed, ...]) -> str:
kinds: Final = tuple(outcome.call.kind for outcome in outcomes)
return str({kind: kinds.count(kind) for kind in _ALL_KINDS if kind in kinds})
def test_c01_upstream_pause_mid_burst_keeps_liveliness_and_lands_every_id_once(
gateway: Gateway,
tmp_path: Path,
record_property: pytest.RecordProperty,
) -> None:
with owned_upstream(tmp_path) as slot, gateway.scenario() as scenario:
upstream: Final = slot.process
assert upstream is not None
deployments: Final = _deployments(scenario, slot.url)
eventually(partial(_every_worker_serves_every_kind, gateway, deployments), lambda served: served, seconds=90)
clients: Final = _clients(gateway, deployments)
calls: Final = _burst_calls(_ALL_KINDS, 5)
observations: Final = _Observations(slot.url)
with ThreadPoolExecutor(max_workers=len(calls)) as pool:
futures: Final = _submit(pool, clients, calls)
eventually(partial(_done_count, futures), lambda done: done >= 1, seconds=60)
upstream.send_signal(signal.SIGSTOP)
try:
done_at_pause: Final = _done_count(futures)
paused_liveliness: Final = tuple(_liveliness(gateway) for _ in range(3))
finally:
upstream.send_signal(signal.SIGCONT)
done_at_resume: Final = _done_count(futures)
outcomes: Final = tuple(future.result(timeout=180) for future in futures)
record_property("c01_burst_size", len(calls))
record_property("c01_done_at_pause", done_at_pause)
record_property("c01_done_at_resume", done_at_resume)
record_property("c01_paused_liveliness_statuses", str(paused_liveliness))
assert paused_liveliness == (200, 200, 200), paused_liveliness
assert done_at_resume < len(calls), done_at_resume
failures: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, _Failed))
assert not failures, failures
completed: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, _Completed))
record_property("c01_completed_by_kind", _kind_counts(completed))
response_ids: Final = tuple(outcome.response_id for outcome in completed)
assert len(set(response_ids)) == len(calls), response_ids
spend_rows: Final = _landed_once(response_ids)
record_property("c01_spend_query", _SPEND_QUERY)
record_property("c01_spend_row_counts", str(tuple(len(rows) for rows in spend_rows)))
markers: Final = tuple(call.marker for call in calls)
eventually(
observations.read,
lambda items: _marker_counts(items, markers) == (1,) * len(markers),
seconds=60,
)
record_property("c01_upstream_marker_counts", str(_marker_counts(observations.items, markers)))
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
def test_c02_worker_sigkill_mid_burst_keeps_serving_and_respawns(
gateway: Gateway,
tmp_path: Path,
record_property: pytest.RecordProperty,
) -> None:
with owned_upstream(tmp_path) as slot, gateway.scenario() as scenario:
deployments: Final = _deployments(scenario, slot.url)
with owned_proxy_process(gateway, tmp_path, {"INTEGRATION_UPSTREAM_URL": slot.url}, workers=2) as owned:
eventually(
partial(_every_worker_serves_every_kind, owned.gateway, deployments), lambda served: served, seconds=90
)
clients: Final = _clients(owned.gateway, deployments)
workers: Final = eventually(
partial(_alive_workers, owned.process), lambda found: len(found) == 2, seconds=60
)
record_property("c02_process_tree_before_kill", _process_tree(owned.process))
victim: Final = workers[0]
calls: Final = _burst_calls(_CHAT_KINDS, 15)
observations: Final = _Observations(slot.url)
with ThreadPoolExecutor(max_workers=len(calls)) as pool:
futures: Final = _submit(pool, clients, calls)
eventually(partial(_done_count, futures), lambda done: done >= 1, seconds=60)
victim.kill()
eventually(partial(_worker_alive, victim), lambda alive: not alive, seconds=10)
done_at_kill: Final = _done_count(futures)
survivor: Final = eventually(
lambda: _attempt(clients, _Call("chat", f"after-kill-{uuid.uuid4().hex}")),
lambda outcome: isinstance(outcome, _Completed),
seconds=30,
)
outcomes: Final = tuple(future.result(timeout=180) for future in futures)
record_property("c02_burst_size", len(calls))
record_property("c02_killed_worker_pid", victim.pid)
record_property("c02_done_at_kill", done_at_kill)
assert isinstance(survivor, _Completed), survivor
completed: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, _Completed))
failed: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, _Failed))
record_property("c02_completed_by_kind", _kind_counts(completed))
record_property("c02_inflight_failure_count", len(failed))
record_property("c02_inflight_failure_errors", str(tuple(outcome.error for outcome in failed)))
assert completed, outcomes
respawned: Final = eventually(
partial(_alive_workers, owned.process),
lambda found: len(found) == 2 and victim.pid not in {worker.pid for worker in found},
seconds=graceful_stop_seconds() + 60,
)
record_property("c02_process_tree_after_respawn", _process_tree(owned.process))
record_property("c02_respawned_worker_pids", str(tuple(worker.pid for worker in respawned)))
after_respawn: Final = eventually(
lambda: _attempt(clients, _Call("chat", f"after-respawn-{uuid.uuid4().hex}")),
lambda outcome: isinstance(outcome, _Completed),
seconds=60,
)
assert isinstance(after_respawn, _Completed), after_respawn
_landed_once((survivor.response_id, after_respawn.response_id))
response_ids: Final = tuple(outcome.response_id for outcome in completed)
assert len(set(response_ids)) == len(completed), response_ids
spend_row_counts: Final = tuple((outcome, len(_spend_rows(outcome.response_id))) for outcome in completed)
landed: Final = tuple(outcome for outcome, count in spend_row_counts if count == 1)
lost: Final = tuple(outcome for outcome, count in spend_row_counts if count == 0)
assert len(landed) + len(lost) == len(completed), spend_row_counts
record_property("c02_spend_query", _SPEND_QUERY)
record_property("c02_burst_ids_landed_once", len(landed))
record_property("c02_burst_ids_lost", len(lost))
record_property("c02_burst_ids_lost_by_kind", _kind_counts(lost))
completed_markers: Final = tuple(outcome.call.marker for outcome in completed)
eventually(
observations.read,
lambda items: _marker_counts(items, completed_markers) == (1,) * len(completed_markers),
seconds=60,
)
failed_markers: Final = tuple(outcome.call.marker for outcome in failed)
record_property(
"c02_failed_calls_seen_upstream",
sum(count == 1 for count in _marker_counts(observations.items, failed_markers)),
)

1
tests/integration/providers/test_databricks_decisions_wire.py
View file

@ -208,7 +208,6 @@ def test_an_openai_format_request_at_v1_decisions_reaches_the_serving_endpoint_a
"model": deployment,
"input": _OPENAI_INPUT,
"questions": [_OPENAI_TIER_QUESTION],
"safety_identifier": "end-user-1",
},
)
assert response.status_code == 200, response.text

586
tests/integration/providers/test_decisions_openai_format_wire.py Normal file
View file

@ -0,0 +1,586 @@
import math
import uuid
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.process import owned_proxy_process
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
from integration.cost_calculation.cost_tracking_case import JsonResponse
from pydantic import JsonValue, TypeAdapter
import litellm
_API_KEY: Final = "synthetic-decisions-key"
_SAFETY_IDENTIFIER: Final = "end-user-7"
_DROPPING_DEPLOYMENT: Final = "decisions-under-litellm-settings-drop-params"
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_INPUT: Final = "Ticket (billing): The export job hangs at 99%"
_FOLLOW_UP: Final = "Customer: still stuck after retrying"
_QUESTIONS: Final[list[JsonValue]] = [
{"type": "predicate", "name": "defect", "instructions": "Is this a defect?"},
{
"type": "choice",
"name": "severity",
"instructions": "How severe is it?",
"choices": [{"value": "low", "description": "cosmetic"}, {"value": "high", "description": "blocks users"}],
},
{
"type": "score",
"name": "confidence",
"instructions": "How sure are you?",
"levels": [{"label": "unsure"}, {"label": "sure"}],
},
]
_SYSTEM_ONE_QUESTIONS: Final[dict[str, JsonValue]] = {
"defect": {"type": "noul", "instructions": "Is this a defect?"},
"severity": {
"type": "choice",
"instructions": "How severe is it?",
"criteria": {"low": "cosmetic", "high": "blocks users"},
},
"confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]},
}
_SYSTEM_ONE_ANSWERS: Final[dict[str, JsonValue]] = {
"defect": {"type": "noul", "noul": 0.93},
"severity": {"type": "choice", "choice": "high", "confidence": 0.8, "probabilities": {"low": 0.2, "high": 0.8}},
"confidence": {
"type": "score",
"score": 1.0,
"confidence": 0.7,
"legend": {"0": "unsure", "1": "sure"},
"probabilities": {"0": 0.3, "1": 0.7},
},
}
_ANSWERS: Final[list[JsonValue]] = [
{"type": "predicate", "name": "defect", "probability": 0.93},
{
"type": "choice",
"name": "severity",
"choice": "high",
"probabilities": [{"value": "low", "probability": 0.2}, {"value": "high", "probability": 0.8}],
"confidence": 0.8,
},
{
"type": "score",
"name": "confidence",
"score": 1.0,
"probabilities": [
{"value": 0, "label": "unsure", "probability": 0.3},
{"value": 1, "label": "sure", "probability": 0.7},
],
"confidence": 0.7,
},
]
_INPUT_TOKENS: Final = 367
_OUTPUT_TOKENS: Final = 3
_SYSTEM_ONE_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": _INPUT_TOKENS, "output_tokens": _OUTPUT_TOKENS}
_USAGE: Final[dict[str, JsonValue]] = {
"input_tokens": _INPUT_TOKENS,
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
"output_tokens": _OUTPUT_TOKENS,
"output_tokens_details": {"reasoning_tokens": 0},
"total_tokens": _INPUT_TOKENS + _OUTPUT_TOKENS,
}
_SDK_QUESTIONS: Final = TypeAdapter(list[dict[str, object]]).validate_python(_QUESTIONS)
_SPEND_QUERY: Final = (
"SELECT spend, status, call_type, model_group, custom_llm_provider, api_base, prompt_tokens, completion_tokens, "
'request_tags FROM "LiteLLM_SpendLogs" WHERE request_id = %s'
)
@dataclass(frozen=True, slots=True)
class _Provider:
name: str
model: str
path: str
body_model: str
api_key: str | None
wraps_result: bool
cost_map_key: str | None
speaks_openai: bool = False
def upstream_body(self) -> dict[str, JsonValue]:
if self.speaks_openai:
return {"model": self.body_model, "input": _INPUT, "questions": _QUESTIONS}
return {"model": self.body_model, "state": _INPUT, "questions": _SYSTEM_ONE_QUESTIONS}
def upstream_reply(self) -> dict[str, JsonValue]:
if self.speaks_openai:
return {"model": self.body_model, "answers": _ANSWERS, "usage": _USAGE}
answer: Final[dict[str, JsonValue]] = {
"model": self.body_model,
"answers": _SYSTEM_ONE_ANSWERS,
"usage": _SYSTEM_ONE_USAGE,
}
return {"result": answer, "success": True} if self.wraps_result else answer
def litellm_response(self) -> dict[str, JsonValue]:
return {"model": self.body_model, "answers": _ANSWERS, "usage": _USAGE}
_PROVIDERS: Final = (
_Provider(
"perplexity",
"perplexity/pplx-decider-v1-27b",
"/v1/decisions",
"pplx-decider-v1-27b",
_API_KEY,
False,
"perplexity/pplx-decider-v1-27b",
),
_Provider("typesafe", "typesafe/jev-1.13.0", "/v1/systemone", "jev-1.13.0", _API_KEY, False, "typesafe/jev-1.13.0"),
_Provider(
"openrouter",
"openrouter/typesafe/jev-1.13",
"/alpha/decisions",
"typesafe/jev-1.13",
_API_KEY,
False,
"openrouter/typesafe/jev-1.13",
),
_Provider(
"strands_decider", "strands_decider/systemone-decider", "/v1/systemone", "systemone-decider", None, False, None
),
_Provider(
"cloudflare",
"cloudflare/clef",
"/ai/run/@cf/cloudflare/clef",
"clef",
_API_KEY,
True,
"cloudflare/@cf/cloudflare/clef",
),
_Provider("hosted_vllm", "hosted_vllm/Qwen/Qwen3-0.6B", "/v1/systemone", "Qwen/Qwen3-0.6B", None, False, None),
_Provider(
"databricks",
"databricks/databricks-openjev-qwen35-4b",
"/databricks-openjev-qwen35-4b/invocations",
"databricks-openjev-qwen35-4b",
_API_KEY,
False,
None,
),
_Provider(
"azure_ai", "azure_ai/decision-1", "/providers/microsoft/v1/systemone", "decision-1", _API_KEY, False, None
),
_Provider("openai", "openai/gpt-6-luna", "/v1/decisions", "gpt-6-luna", _API_KEY, False, "gpt-6-luna", True),
)
_SYSTEM_ONE_PROVIDERS: Final = tuple(provider for provider in _PROVIDERS if not provider.speaks_openai)
_PERPLEXITY: Final = _PROVIDERS[0]
_TYPESAFE: Final = _PROVIDERS[1]
_OPENAI: Final = next(provider for provider in _PROVIDERS if provider.speaks_openai)
_PREDICATE: Final[dict[str, JsonValue]] = {"type": "predicate", "name": "q", "instructions": "Is it?"}
_INVALID_BODIES: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = (
("missing questions", {"input": _INPUT}),
("missing input", {"questions": _QUESTIONS}),
("numeric input", {"input": 5, "questions": _QUESTIONS}),
("assistant message input", {"input": [{"role": "assistant", "content": "hi"}], "questions": _QUESTIONS}),
("empty questions", {"input": _INPUT, "questions": []}),
("questions as a map", {"input": _INPUT, "questions": {"q": _PREDICATE}}),
("predicate without instructions", {"input": _INPUT, "questions": [{"type": "predicate", "name": "q"}]}),
(
"choice with one choice",
{"input": _INPUT, "questions": [{**_PREDICATE, "type": "choice", "choices": [{"value": "only"}]}]},
),
(
"score with one level",
{"input": _INPUT, "questions": [{**_PREDICATE, "type": "score", "levels": [{"label": "only"}]}]},
),
("unknown question type", {"input": _INPUT, "questions": [{**_PREDICATE, "type": "ranking"}]}),
("numeric safety_identifier", {"input": _INPUT, "questions": _QUESTIONS, "safety_identifier": 7}),
("list safety_identifier", {"input": _INPUT, "questions": _QUESTIONS, "safety_identifier": [_SAFETY_IDENTIFIER]}),
)
_IMAGE_INPUT: Final[list[JsonValue]] = [
{
"role": "user",
"content": [
{"type": "input_text", "text": _INPUT},
{"type": "input_image", "image_url": "data:image/png;base64,iVBORw0KGgo=", "detail": "auto"},
],
}
]
_OPENAI_IMAGE_MESSAGES: Final[list[JsonValue]] = [
{
"role": "user",
"type": "message",
"content": [
{"type": "input_text", "text": _INPUT},
{"type": "input_image", "image_url": "data:image/png;base64,iVBORw0KGgo=", "detail": "auto"},
],
}
]
_MESSAGE_LIST_INPUT: Final[list[JsonValue]] = [
{"role": "user", "content": _INPUT},
{"role": "user", "content": [{"type": "input_text", "text": _FOLLOW_UP}]},
]
def _response_cost(response: httpx.Response) -> float:
return float(response.headers["x-litellm-response-cost"]) if "x-litellm-response-cost" in response.headers else 0.0
def _provider_id(provider: _Provider) -> str:
return provider.name
def _number(value: JsonValue) -> float:
assert isinstance(value, (int, float)) and not isinstance(value, bool), value
return float(value)
def _expected_spend(cost_map_key: str | None) -> float:
if cost_map_key is None:
return 0.0
cost_map: Final = _JSON_OBJECT.validate_json(Path("model_prices_and_context_window.json").read_bytes())
prices: Final = object_value(cost_map[cost_map_key])
return _INPUT_TOKENS * _number(prices["input_cost_per_token"]) + _OUTPUT_TOKENS * _number(
prices["output_cost_per_token"]
)
def _register(scenario: Scenario, body: dict[str, JsonValue], *, status: int = 200) -> ScenarioHandle:
handle: Final = register_scenario(
f"decisions-{uuid.uuid4().hex[:12]}", JsonResponse(content_type="application/json", body=body, status=status)
)
scenario.cleanups.callback(delete_scenario, handle)
return handle
def _deployment(scenario: Scenario, handle: ScenarioHandle, provider: _Provider, *, drop_params: bool = False) -> str:
if drop_params:
return scenario.model(
model=provider.model, api_base=handle.api_base(), api_key=provider.api_key, drop_params=True
)
return scenario.model(model=provider.model, api_base=handle.api_base(), api_key=provider.api_key)
def _decide(gateway: Gateway, model: str, *, route: str = "/v1/decisions", **extra: JsonValue) -> httpx.Response:
return gateway.request("POST", route, {"model": model, "input": _INPUT, "questions": _QUESTIONS, **extra})
def _decide_in_system_one_format(gateway: Gateway, model: str, **extra: JsonValue) -> httpx.Response:
return gateway.request(
"POST", "/v1/systemone", {"model": model, "state": _INPUT, "questions": _SYSTEM_ONE_QUESTIONS, **extra}
)
def _observed_requests(gateway: Gateway) -> tuple[dict[str, JsonValue], ...]:
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
observations: Final = _JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"]
assert isinstance(observations, list), observations
return tuple(map(object_value, observations))
def _calls_to(requests: Sequence[dict[str, JsonValue]], handle: ScenarioHandle) -> list[dict[str, JsonValue]]:
return [request for request in requests if string_value(request["path"]).startswith(f"/{handle.scenario_id}/")]
def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, JsonValue]]:
return _calls_to(_observed_requests(gateway), handle)
def _spend_row(call_id: str) -> dict[str, JsonValue]:
rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70)
return rows[0]
def _assert_refused_as_invalid(gateway: Gateway, model: str, label: str, body: dict[str, JsonValue]) -> None:
response: Final = gateway.request("POST", "/v1/decisions", {"model": model, **body})
assert response.status_code == 400, (label, response.text)
assert "Invalid Decisions request" in response.text, (label, response.text)
def _assert_refused_for_safety_identifier(response: httpx.Response) -> None:
assert response.status_code == 400, response.text
assert "safety_identifier" in response.text and "drop_params" in response.text, response.text
def _drop_params_config(directory: Path, api_base: str) -> Path:
base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
config: Final = {
**base,
"litellm_settings": {**object_value(base["litellm_settings"]), "drop_params": True},
"model_list": [
{
"model_name": _DROPPING_DEPLOYMENT,
"litellm_params": {"model": _PERPLEXITY.model, "api_base": api_base, "api_key": _API_KEY},
}
],
}
path: Final = directory / "decisions-drop-params.yaml"
path.write_text(yaml.safe_dump(config))
return path
@pytest.mark.parametrize("provider", _PROVIDERS, ids=_provider_id)
def test_each_provider_gets_its_own_path_key_and_body_and_is_billed_from_the_cost_map(
gateway: Gateway, provider: _Provider
) -> None:
expected_spend: Final = _expected_spend(provider.cost_map_key)
with gateway.scenario() as scenario:
handle: Final = _register(scenario, provider.upstream_reply())
model: Final = _deployment(scenario, handle, provider)
response: Final = _decide(gateway, model)
assert response.status_code == 200, response.text
assert response.json() == provider.litellm_response()
assert response.headers["x-litellm-model-group"] == model
assert math.isclose(_response_cost(response), expected_spend, rel_tol=1e-9)
(call,) = _upstream_calls(gateway, handle)
assert call["path"] == f"/{handle.scenario_id}{provider.path}"
assert call["authorization"] == (f"Bearer {provider.api_key}" if provider.api_key else "")
assert call["body"] == provider.upstream_body()
row: Final = _spend_row(response.headers["x-litellm-call-id"])
assert (
row["status"],
row["call_type"],
row["custom_llm_provider"],
row["model_group"],
row["api_base"],
row["prompt_tokens"],
row["completion_tokens"],
) == (
"success",
"adecisions",
provider.name,
model,
f"{handle.api_base()}{provider.path}",
_INPUT_TOKENS,
_OUTPUT_TOKENS,
)
assert math.isclose(_number(row["spend"]), expected_spend, rel_tol=1e-9), row
def test_the_unversioned_alias_serves_the_same_request(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _PERPLEXITY.upstream_reply())
model: Final = _deployment(scenario, handle, _PERPLEXITY)
response: Final = _decide(gateway, model, route="/decisions")
assert response.status_code == 200, response.text
assert response.json() == _PERPLEXITY.litellm_response()
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == _PERPLEXITY.upstream_body()
def test_repeated_identical_requests_each_reach_the_upstream_and_are_each_billed(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _PERPLEXITY.upstream_reply())
model: Final = _deployment(scenario, handle, _PERPLEXITY)
responses: Final = tuple(_decide(gateway, model) for _ in range(2))
assert [response.status_code for response in responses] == [200, 200], [r.text for r in responses]
call_ids: Final = tuple(response.headers["x-litellm-call-id"] for response in responses)
assert len(set(call_ids)) == 2, call_ids
assert len(_upstream_calls(gateway, handle)) == 2
for call_id in call_ids:
assert _spend_row(call_id)["status"] == "success"
async def test_sdk_sync_and_async_clients_send_the_same_request(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _TYPESAFE.upstream_reply())
synchronous: Final = litellm.decisions(
model=_TYPESAFE.model, input=_INPUT, questions=_SDK_QUESTIONS, api_base=handle.api_base(), api_key=_API_KEY
)
asynchronous: Final = await litellm.adecisions(
model=_TYPESAFE.model, input=_INPUT, questions=_SDK_QUESTIONS, api_base=handle.api_base(), api_key=_API_KEY
)
for response in (synchronous, asynchronous):
assert response.model_dump(mode="json") == _TYPESAFE.litellm_response()
calls: Final = _upstream_calls(gateway, handle)
assert len(calls) == 2, calls
for call in calls:
assert call["path"] == f"/{handle.scenario_id}{_TYPESAFE.path}"
assert call["authorization"] == f"Bearer {_API_KEY}"
assert call["body"] == _TYPESAFE.upstream_body()
def test_gateway_only_fields_stay_at_the_gateway_and_tags_reach_the_spend_log(gateway: Gateway) -> None:
tag: Final = f"decisions-openai-format-{uuid.uuid4().hex[:8]}"
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _PERPLEXITY.upstream_reply())
model: Final = _deployment(scenario, handle, _PERPLEXITY)
response: Final = _decide(
gateway, model, user="auditor", num_retries=0, temperature=0.2, metadata={"tags": [tag]}
)
assert response.status_code == 200, response.text
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == _PERPLEXITY.upstream_body()
row: Final = _spend_row(response.headers["x-litellm-call-id"])
tags: Final = row["request_tags"]
assert isinstance(tags, list) and tag in tags, row
def test_invalid_bodies_are_refused_at_the_gateway_without_an_upstream_call(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _PERPLEXITY.upstream_reply())
model: Final = _deployment(scenario, handle, _PERPLEXITY)
for label, body in _INVALID_BODIES:
_assert_refused_as_invalid(gateway, model, label, body)
assert _upstream_calls(gateway, handle) == []
@pytest.mark.parametrize("provider", _SYSTEM_ONE_PROVIDERS, ids=_provider_id)
def test_system_one_providers_refuse_images_at_the_gateway_without_an_upstream_call(
gateway: Gateway, provider: _Provider
) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, provider.upstream_reply())
model: Final = _deployment(scenario, handle, provider)
response: Final = _decide(gateway, model, input=_IMAGE_INPUT)
assert response.status_code == 400, response.text
assert "input_image" in response.text
assert _upstream_calls(gateway, handle) == []
def test_openai_forwards_image_input_in_its_own_message_shape(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _OPENAI.upstream_reply())
model: Final = _deployment(scenario, handle, _OPENAI)
response: Final = _decide(gateway, model, input=_IMAGE_INPUT)
assert response.status_code == 200, response.text
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == {"model": _OPENAI.body_model, "input": _OPENAI_IMAGE_MESSAGES, "questions": _QUESTIONS}
def test_a_message_list_input_reaches_a_system_one_provider_as_its_flattened_text(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _PERPLEXITY.upstream_reply())
model: Final = _deployment(scenario, handle, _PERPLEXITY)
response: Final = _decide(gateway, model, input=_MESSAGE_LIST_INPUT)
assert response.status_code == 200, response.text
assert response.json() == _PERPLEXITY.litellm_response()
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == {
"model": _PERPLEXITY.body_model,
"state": f"{_INPUT}\n\n{_FOLLOW_UP}",
"questions": _SYSTEM_ONE_QUESTIONS,
}
row: Final = _spend_row(response.headers["x-litellm-call-id"])
assert (row["status"], row["call_type"], row["prompt_tokens"], row["completion_tokens"]) == (
"success",
"adecisions",
_INPUT_TOKENS,
_OUTPUT_TOKENS,
)
@pytest.mark.parametrize("provider", _SYSTEM_ONE_PROVIDERS, ids=_provider_id)
def test_safety_identifier_is_refused_by_system_one_providers_unless_the_deployment_drops_params(
gateway: Gateway, provider: _Provider
) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, provider.upstream_reply())
strict: Final = _deployment(scenario, handle, provider)
refused: Final = _decide(gateway, strict, safety_identifier=_SAFETY_IDENTIFIER)
_assert_refused_for_safety_identifier(refused)
assert _upstream_calls(gateway, handle) == []
dropping: Final = _deployment(scenario, handle, provider, drop_params=True)
accepted: Final = _decide(gateway, dropping, safety_identifier=_SAFETY_IDENTIFIER)
assert accepted.status_code == 200, accepted.text
assert accepted.json() == provider.litellm_response()
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == provider.upstream_body()
row: Final = _spend_row(accepted.headers["x-litellm-call-id"])
assert (row["status"], row["call_type"], row["model_group"]) == ("success", "adecisions", dropping)
@pytest.mark.parametrize("value", ("", "x" * 5000), ids=("empty", "5kb"))
def test_every_string_safety_identifier_is_refused_or_dropped_like_the_usual_one(gateway: Gateway, value: str) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _PERPLEXITY.upstream_reply())
strict: Final = _deployment(scenario, handle, _PERPLEXITY)
_assert_refused_for_safety_identifier(_decide(gateway, strict, safety_identifier=value))
dropping: Final = _deployment(scenario, handle, _PERPLEXITY, drop_params=True)
accepted: Final = _decide(gateway, dropping, safety_identifier=value)
assert accepted.status_code == 200, accepted.text
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == _PERPLEXITY.upstream_body()
def test_a_request_body_drop_params_drops_the_safety_identifier_like_chat(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _PERPLEXITY.upstream_reply())
strict: Final = _deployment(scenario, handle, _PERPLEXITY)
response: Final = _decide(gateway, strict, safety_identifier=_SAFETY_IDENTIFIER, drop_params=True)
assert response.status_code == 200, response.text
assert response.json() == _PERPLEXITY.litellm_response()
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == _PERPLEXITY.upstream_body()
def test_openai_keeps_the_safety_identifier_on_the_wire_without_drop_params(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _OPENAI.upstream_reply())
model: Final = _deployment(scenario, handle, _OPENAI)
response: Final = _decide(gateway, model, safety_identifier=_SAFETY_IDENTIFIER)
assert response.status_code == 200, response.text
assert response.json() == _OPENAI.litellm_response()
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == {**_OPENAI.upstream_body(), "safety_identifier": _SAFETY_IDENTIFIER}
def test_a_system_one_format_safety_identifier_is_refused_unless_the_deployment_drops_params(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _PERPLEXITY.upstream_reply())
strict: Final = _deployment(scenario, handle, _PERPLEXITY)
refused: Final = _decide_in_system_one_format(gateway, strict, safety_identifier=_SAFETY_IDENTIFIER)
_assert_refused_for_safety_identifier(refused)
assert _upstream_calls(gateway, handle) == []
dropping: Final = _deployment(scenario, handle, _PERPLEXITY, drop_params=True)
accepted: Final = _decide_in_system_one_format(gateway, dropping, safety_identifier=_SAFETY_IDENTIFIER)
assert accepted.status_code == 200, accepted.text
assert accepted.json()["answers"] == _SYSTEM_ONE_ANSWERS
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == _PERPLEXITY.upstream_body()
def test_openai_keeps_a_system_one_format_safety_identifier_on_the_wire(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _OPENAI.upstream_reply())
model: Final = _deployment(scenario, handle, _OPENAI)
response: Final = _decide_in_system_one_format(gateway, model, safety_identifier=_SAFETY_IDENTIFIER)
assert response.status_code == 200, response.text
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == {**_OPENAI.upstream_body(), "safety_identifier": _SAFETY_IDENTIFIER}
@pytest.mark.parametrize("decide", (_decide, _decide_in_system_one_format), ids=("openai_format", "system_one_format"))
@pytest.mark.parametrize("value", (7, [_SAFETY_IDENTIFIER]), ids=("numeric", "list"))
def test_a_non_string_safety_identifier_is_refused_as_invalid_unless_the_deployment_drops_params(
gateway: Gateway, value: JsonValue, decide: Callable[..., httpx.Response]
) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _OPENAI.upstream_reply())
strict: Final = _deployment(scenario, handle, _OPENAI)
refused: Final = decide(gateway, strict, safety_identifier=value)
assert refused.status_code == 400, refused.text
assert "Invalid Decisions request" in refused.text, refused.text
assert _upstream_calls(gateway, handle) == []
dropping: Final = _deployment(scenario, handle, _OPENAI, drop_params=True)
accepted: Final = decide(gateway, dropping, safety_identifier=value)
assert accepted.status_code == 200, accepted.text
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == _OPENAI.upstream_body()
def test_litellm_settings_drop_params_drops_the_safety_identifier_for_a_strict_deployment(
gateway: Gateway, tmp_path: Path
) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _PERPLEXITY.upstream_reply())
config: Final = _drop_params_config(tmp_path, handle.api_base())
with owned_proxy_process(gateway, tmp_path, {}, config=config) as owned:
response: Final = _decide(owned.gateway, _DROPPING_DEPLOYMENT, safety_identifier=_SAFETY_IDENTIFIER)
assert response.status_code == 200, response.text
assert response.json() == _PERPLEXITY.litellm_response()
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == _PERPLEXITY.upstream_body()

29
tests/integration/providers/test_realtime_transcription_guardrail_session_update.py
View file

@ -1139,7 +1139,7 @@ def test_muse_session_without_intent_still_delivers_the_transcript_and_the_verdi
assert _muse_handshake(observed)["model"] == MUSE_MODEL, _muse_handshake(observed)
def test_without_a_guardrail_the_proxy_sends_no_session_update_of_its_own(gateway: Gateway) -> None:
def test_without_a_guardrail_the_proxy_sends_no_session_update_or_response_create_of_its_own(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
transcription: Final = _scripted(scenario, _transcription_scenario(BLOCKED_TRANSCRIPT))
voice: Final = _scripted(scenario, _voice_scenario(BLOCKED_TRANSCRIPT))
@ -1147,22 +1147,29 @@ def test_without_a_guardrail_the_proxy_sends_no_session_update_of_its_own(gatewa
transcribe_model: Final = _openai_deployment(scenario, transcription.scenario_id)
voice_model: Final = _openai_deployment(scenario, voice.scenario_id, model=VOICE_MODEL)
transcribed: Final = _transcribe(_ws_base(_proxy_url()), f"model={transcribe_model}&{TRANSCRIPTION_QUERY}", key)
talked: Final = _talk(_ws_base(_proxy_url()), f"model={voice_model}", key, _voice_frames(VOICE_UPDATE))
talked: Final = _talk_without_a_reply(_ws_base(_proxy_url()), f"model={voice_model}", key)
observed: Final = _observed_all(gateway.upstream_url)
transcription_observed: Final = _belonging(observed, transcription.scenario_id)
voice_observed: Final = _belonging(observed, voice.scenario_id)
_assert_transcription_left_alone(
transcribed, transcription_observed, BLOCKED_TRANSCRIPT, update=GA_TRANSCRIPTION_UPDATE
)
assert talked.types == ("session.created", "session.updated", TRANSCRIPT_COMPLETED, RESPONSE_DONE), talked
assert talked.errors == (), talked
assert _session_updates(voice_observed) == (VOICE_UPDATE,), _session_updates(voice_observed)
_assert_voice_session_ungated(talked, _belonging(observed, voice.scenario_id))
def _talk_without_a_reply(ws_base: str, query: str, key: str) -> Session:
return _drive(
ws_base,
query,
key,
(_step(VOICE_UPDATE_FRAME, COMMIT, until=TRANSCRIPT_COMPLETED), _step(until=RESPONSE_DONE)),
)
def _assert_voice_session_ungated(session: Session, observed: tuple[dict[str, JsonValue], ...]) -> None:
assert session.types == ("session.created", "session.updated", TRANSCRIPT_COMPLETED, RESPONSE_DONE), session
assert session.types == ("session.created", "session.updated", TRANSCRIPT_COMPLETED, "timeout"), session
assert session.transcripts == (BLOCKED_TRANSCRIPT,), session
assert session.errors == (), session
assert _sent_types(observed) == ("session.update", "input_audio_buffer.commit"), _sent(observed)
assert _session_updates(observed) == (VOICE_UPDATE,), _session_updates(observed)
@ -1180,9 +1187,7 @@ def test_key_opted_out_of_the_transcript_guardrail_is_not_gated_by_the_prompt_gu
handle: Final = _scripted(scenario, _voice_scenario(BLOCKED_TRANSCRIPT))
key: Final = scenario.key(metadata={"opted_out_global_guardrails": [TRANSCRIPT_GUARDRAIL]})
model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL)
session: Final = _talk(
_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, _voice_frames(VOICE_UPDATE)
)
session: Final = _talk_without_a_reply(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key)
_assert_voice_session_ungated(session, _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id))
@ -1199,9 +1204,7 @@ def test_team_opted_out_of_the_transcript_guardrail_is_not_gated(guardrail_proxy
team: Final = scenario.team(metadata={"opted_out_global_guardrails": [TRANSCRIPT_GUARDRAIL]})
key: Final = scenario.key(team_id=team)
model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL)
session: Final = _talk(
_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, _voice_frames(VOICE_UPDATE)
)
session: Final = _talk_without_a_reply(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key)
_assert_voice_session_ungated(session, _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id))

1004
tests/integration/providers/test_realtime_unguarded_voice_turns.py Normal file
View file

File diff suppressed because it is too large Load diff

827
tests/integration/providers/test_strands_decider_agentcore_wire.py Normal file
View file

@ -0,0 +1,827 @@
import hashlib
import json
import os
import re
import threading
import time
import uuid
from collections.abc import Callable, Iterator, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, replace
from functools import partial
from pathlib import Path
from typing import Final
from urllib.parse import quote, unquote
import httpx
import pytest
import yaml
from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, gateway_from_environment
from integration._support.database import read_rows
from integration._support.forward_proxy import ForwardProxy, tunnelling_forward_proxy
from integration._support.process import owned_proxy
from integration._support.sigv4 import encoded_path, signature
from integration._support.tls import server_context, write_self_signed_cert
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.llms.strands_decider.decisions.transformation import StrandsDeciderDecisionsConfig
_ACCOUNT: Final = "123456789012"
_ACCESS_KEY: Final = "AKIAAGENTCOREAUDIT01"
_SECRET_KEY: Final = "synthetic-agentcore-audit-secret"
_SIGNING: Final[Mapping[str, JsonValue]] = {"aws_access_key_id": _ACCESS_KEY, "aws_secret_access_key": _SECRET_KEY}
_JWT: Final = "synthetic-agentcore-jwt"
_ENV_JWT: Final = "synthetic-agentcore-env-jwt"
_US_EAST: Final = "us-east-1"
_EU_WEST: Final = "eu-west-1"
_MODEL: Final = "strands_decider/systemone-decider"
_BODY_MODEL: Final = "systemone-decider"
_SESSION_HEADER: Final = "x-amzn-bedrock-agentcore-runtime-session-id"
_CONFIGURED_SESSION: Final = "agentcore-audit-configured-session-0000000001"
_STATE: Final[dict[str, JsonValue]] = {"ticket": "Refund requested twice for one order", "queue": "payments"}
_QUESTIONS: Final[dict[str, JsonValue]] = {
"duplicate": {"type": "noul", "instructions": "Is this a duplicate refund request?"},
"urgency": {"type": "choice", "criteria": {"low": "can wait", "high": "money at risk"}},
}
_SDK_QUESTIONS: Final = TypeAdapter(dict[str, dict[str, object]]).validate_python(_QUESTIONS)
_ANSWERS: Final[dict[str, JsonValue]] = {
"duplicate": {"type": "noul", "noul": 0.88},
"urgency": {"type": "choice", "choice": "high", "confidence": 0.75, "probabilities": {"low": 0.25, "high": 0.75}},
}
_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 211, "output_tokens": 2}
_UPSTREAM_BODY: Final[dict[str, JsonValue]] = {"model": _BODY_MODEL, "state": _STATE, "questions": _QUESTIONS}
_INPUT: Final = "Refund requested twice for one order"
_OPENAI_QUESTIONS: Final[list[JsonValue]] = [
{"type": "predicate", "name": "duplicate", "instructions": "Is this a duplicate refund request?"},
{
"type": "choice",
"name": "urgency",
"instructions": "How urgent is it?",
"choices": [{"value": "low", "description": "can wait"}, {"value": "high", "description": "money at risk"}],
},
]
_OPENAI_UPSTREAM_BODY: Final[dict[str, JsonValue]] = {
"model": _BODY_MODEL,
"state": _INPUT,
"questions": {
"duplicate": {"type": "noul", "instructions": "Is this a duplicate refund request?"},
"urgency": {
"type": "choice",
"instructions": "How urgent is it?",
"criteria": {"low": "can wait", "high": "money at risk"},
},
},
}
_OPENAI_ANSWER: Final[dict[str, JsonValue]] = {
"model": _BODY_MODEL,
"answers": [
{"type": "predicate", "name": "duplicate", "probability": 0.88},
{
"type": "choice",
"name": "urgency",
"choice": "high",
"probabilities": [{"value": "low", "probability": 0.25}, {"value": "high", "probability": 0.75}],
"confidence": 0.75,
},
],
"usage": {
"input_tokens": 211,
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
"output_tokens": 2,
"output_tokens_details": {"reasoning_tokens": 0},
"total_tokens": 213,
},
}
@dataclass(frozen=True, slots=True)
class _Route:
path: str
body: Mapping[str, JsonValue]
upstream: Mapping[str, JsonValue]
answer: Mapping[str, JsonValue]
_SYSTEMONE: Final = _Route(
"/v1/systemone",
{"state": _STATE, "questions": _QUESTIONS},
_UPSTREAM_BODY,
{"model": _BODY_MODEL, "answers": _ANSWERS, "usage": _USAGE},
)
_OPENAI: Final = _Route(
"/v1/decisions", {"input": _INPUT, "questions": _OPENAI_QUESTIONS}, _OPENAI_UPSTREAM_BODY, _OPENAI_ANSWER
)
_ROUTES: Final = (_SYSTEMONE, _OPENAI, replace(_SYSTEMONE, path="/systemone"), replace(_OPENAI, path="/decisions"))
_SPEND_QUERY: Final = (
"SELECT status, call_type, custom_llm_provider, model_group, api_base, prompt_tokens, completion_tokens "
'FROM "LiteLLM_SpendLogs" WHERE request_id = %s'
)
_RUNTIME_TARGET: Final = re.compile(
r"\A/runtimes/arn:aws:bedrock-agentcore:[a-z0-9-]+:[0-9]{12}:runtime/(?P<runtime>[A-Za-z0-9_-]+)/invocations\Z"
)
_ERROR_STATUS: Final = {"bad_request": 400, "invalid_request": 400, "loading": 503, "inference_error": 500}
_HTTP_ERRORS: Final = {
403: "AccessDeniedException",
404: "ResourceNotFoundException",
429: "ThrottlingException",
}
_MALFORMED_SESSIONS: Final[Mapping[str, JsonValue]] = {"int": 1234567, "list": ["a" * 40], "empty": ""}
_LONG_SESSION: Final = "s" * 5120
_NEAR_MISS_ARNS: Final = {
"endpoint": f"arn:aws:bedrock-agentcore:us-east-1:{_ACCOUNT}:runtime/audit_qualified/runtime-endpoint/DEFAULT",
"account": "arn:aws:bedrock-agentcore:us-east-1:12345:runtime/audit_short_account",
"service": f"arn:aws:bedrock:us-east-1:{_ACCOUNT}:runtime/audit_other_service",
"host": f"arn:aws:bedrock-agentcore:us-east-1.example.com:{_ACCOUNT}:runtime/audit_crafted_host",
}
_DEPLOYMENT_HEADERS: Final[dict[str, JsonValue]] = {
"x-amz-target": "AmazonBedrockAgentCore.Forged",
"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": "deployment-forged-session-00000000000001",
"Host": "bedrock-agentcore.us-west-2.amazonaws.com",
"x-audit-trace": "deployment-trace",
}
_SIGNED: Final = "agentcore-audit-signed"
_REGION: Final = "agentcore-audit-region"
_SESSION: Final = "agentcore-audit-session"
_LONG: Final = "agentcore-audit-session-long"
_BEARER: Final = "agentcore-audit-bearer"
_HEADERS: Final = "agentcore-audit-headers"
_OPT_IN: Final = "agentcore-audit-opt-in"
_HEALTH: Final = "agentcore-audit-health"
_DROPPED: Final = "agentcore-audit-dropped"
_OUTAGE: Final = "agentcore-audit-outage"
_ENV: Final = "agentcore-audit-env"
_UNSIGNED: Final = "agentcore-audit-unsigned"
_SCRUBBED_PREFIXES: Final = ("AWS_", "STRANDS_DECIDER_")
_SCRUBBED_NAMES: Final = frozenset(
{"HTTPS_PROXY", "HTTP_PROXY", "ALL_PROXY", "NO_PROXY", "SSL_VERIFY", "SSL_CERT_FILE", "REQUESTS_CA_BUNDLE"}
)
_load_yaml: Final[Callable[[str], object]] = yaml.safe_load
_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_OWNED_PROXY_CELL_SECONDS: Final = 2 * max(30.0, float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS") or 70)) + 120
pytestmark: Final = pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
def _host(region: str) -> str:
return f"bedrock-agentcore.{region}.amazonaws.com"
def _arn(runtime: str, region: str = _US_EAST) -> str:
return f"arn:aws:bedrock-agentcore:{region}:{_ACCOUNT}:runtime/{runtime}"
def _invocations_url(runtime: str, region: str = _US_EAST) -> str:
return f"https://{_host(region)}/runtimes/{quote(_arn(runtime, region), safe='')}/invocations"
def _default_session(runtime: str, region: str = _US_EAST) -> str:
return f"litellm-decider-{hashlib.sha256(_arn(runtime, region).encode()).hexdigest()}"
def _runtime_of(request: Request) -> str | None:
match: Final = _RUNTIME_TARGET.match(unquote(request.target))
return None if match is None else match["runtime"]
def _json_reply(body: Mapping[str, JsonValue], status: int = 200, **headers: str) -> Reply:
return Reply(status=status, body=json.dumps(body).encode(), headers=headers)
def _respond(outage: threading.Event, request: Request) -> Reply:
runtime: Final = _runtime_of(request) or ""
if runtime == "audit_dropped" or (runtime == "audit_outage" and outage.is_set()):
return Reply(drop_connection=True)
if runtime.startswith("audit_error_"):
code: Final = runtime.removeprefix("audit_error_")
return _json_reply({"error": {"code": code, "message": f"runtime refused with {code}"}})
if runtime.startswith("audit_status_"):
status: Final = int(runtime.removeprefix("audit_status_"))
return _json_reply(
{"message": f"scripted {_HTTP_ERRORS[status]}"}, status, **{"x-amzn-errortype": _HTTP_ERRORS[status]}
)
if not runtime:
return _json_reply({"message": f"unknown target {request.target}"}, 404)
return _json_reply({"model": _BODY_MODEL, "answers": _ANSWERS, "usage": _USAGE})
@dataclass(frozen=True, slots=True)
class _Runtime:
wire: Wire
tunnel: ForwardProxy
cert: Path
outage: threading.Event
def reset(self) -> None:
self.outage.clear()
self.wire.drain()
self.tunnel.drain()
def calls(self) -> tuple[Request, ...]:
return self.wire.drain()
def _calls_to(calls: Sequence[Request], runtime: str) -> tuple[Request, ...]:
return tuple(call for call in calls if _runtime_of(call) == runtime)
def _deployment(name: str, **litellm_params: JsonValue) -> dict[str, JsonValue]:
return {
"model_name": name,
"litellm_params": {"model": _MODEL, **litellm_params},
"model_info": {"id": name, "mode": "evaluation"},
}
def _signed_deployment(name: str, runtime: str, region: str = _US_EAST, **extra: JsonValue) -> dict[str, JsonValue]:
return _deployment(name, api_base=_arn(runtime, region), **_SIGNING, **extra)
def _deployments() -> list[dict[str, JsonValue]]:
return [
_signed_deployment(_SIGNED, "audit_signed", aws_region_name=_US_EAST),
_signed_deployment(_REGION, "audit_region", _EU_WEST, aws_region_name=_US_EAST),
_signed_deployment(_SESSION, "audit_session", agentcore_runtime_session_id=_CONFIGURED_SESSION),
_signed_deployment(_LONG, "audit_session_long", agentcore_runtime_session_id=_LONG_SESSION),
_deployment(_BEARER, api_base=_arn("audit_bearer"), api_key=_JWT),
_signed_deployment(_HEADERS, "audit_headers", extra_headers=_DEPLOYMENT_HEADERS),
_signed_deployment(
_OPT_IN, "audit_opt_in", configurable_clientside_auth_params=["agentcore_runtime_session_id"]
),
_signed_deployment(_HEALTH, "audit_health"),
_signed_deployment(_DROPPED, "audit_dropped"),
_signed_deployment(_OUTAGE, "audit_outage"),
_deployment(_ENV, **_SIGNING),
_deployment(_UNSIGNED, api_base=_arn("audit_unsigned")),
*(_signed_deployment(f"agentcore-audit-error-{code}", f"audit_error_{code}") for code in _ERROR_STATUS),
*(_signed_deployment(f"agentcore-audit-status-{status}", f"audit_status_{status}") for status in _HTTP_ERRORS),
*(
_signed_deployment(
f"agentcore-audit-session-{label}", f"audit_session_{label}", agentcore_runtime_session_id=value
)
for label, value in _MALFORMED_SESSIONS.items()
),
*(
_deployment(f"agentcore-audit-near-miss-{label}", api_base=value, **_SIGNING)
for label, value in _NEAR_MISS_ARNS.items()
),
]
def _write_config(directory: Path) -> Path:
base: Final = JSON_OBJECT.validate_python(_load_yaml(Path("tests/integration/proxy_config.yaml").read_text()))
configuration: Final = {
**base,
"model_list": _deployments(),
"router_settings": {"disable_cooldowns": True, "num_retries": 0},
}
path: Final = directory / "agentcore.yaml"
path.write_text(yaml.safe_dump(configuration))
return path
def _scrubbed() -> tuple[str, ...]:
return tuple(
name for name in os.environ if name.upper().startswith(_SCRUBBED_PREFIXES) or name.upper() in _SCRUBBED_NAMES
)
def _isolated_aws(directory: Path) -> dict[str, str]:
return {
"AWS_EC2_METADATA_DISABLED": "true",
"AWS_SHARED_CREDENTIALS_FILE": str(directory / "absent-credentials"),
"AWS_CONFIG_FILE": str(directory / "absent-config"),
}
def _proxy_environment(runtime: _Runtime, directory: Path) -> dict[str, str]:
return {
"HTTPS_PROXY": runtime.tunnel.url,
"NO_PROXY": "127.0.0.1,localhost",
"SSL_CERT_FILE": str(runtime.cert),
"STRANDS_DECIDER_API_BASE": _arn("audit_env"),
**_isolated_aws(directory),
}
@pytest.fixture(scope="module")
def module_gateway() -> Iterator[Gateway]:
with gateway_from_environment() as value:
yield value
@pytest.fixture(scope="module")
def runtime(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Runtime]:
cert, key = write_self_signed_cert(tmp_path_factory.mktemp("agentcore-cert"), (_host(_US_EAST), _host(_EU_WEST)))
outage: Final = threading.Event()
with wire_server(partial(_respond, outage), tls=server_context(cert, key)) as wire:
with tunnelling_forward_proxy(int(wire.url.rsplit(":", 1)[1])) as tunnel:
yield _Runtime(wire, tunnel, cert, outage)
@pytest.fixture(scope="module")
def agentcore(
module_gateway: Gateway, runtime: _Runtime, tmp_path_factory: pytest.TempPathFactory
) -> Iterator[Gateway]:
directory: Final = tmp_path_factory.mktemp("agentcore-proxy")
with owned_proxy(
module_gateway,
directory,
_proxy_environment(runtime, directory),
config=_write_config(directory),
remove_environment=_scrubbed(),
workers=2,
) as candidate:
yield candidate
def _decide(gateway: Gateway, model: str, route: _Route = _SYSTEMONE, **extra: JsonValue) -> httpx.Response:
return gateway.request("POST", route.path, {"model": model, **route.body, **extra})
def _authorization_fields(request: Request) -> dict[str, str]:
authorization: Final = request.headers.get("authorization", "")
assert authorization.startswith("AWS4-HMAC-SHA256 "), dict(request.headers)
return dict(part.split("=", 1) for part in authorization.removeprefix("AWS4-HMAC-SHA256 ").split(", "))
def _assert_signed(request: Request, region: str = _US_EAST) -> None:
fields: Final = _authorization_fields(request)
access_key, scope = fields["Credential"].split("/", 1)
assert access_key == _ACCESS_KEY, fields
assert scope == f"{request.headers['x-amz-date'][:8]}/{region}/bedrock-agentcore/aws4_request", fields
signed: Final = fields["SignedHeaders"]
assert {"host", "x-amz-date", _SESSION_HEADER}.issubset(signed.split(";")), signed
expected: Final = signature(
"POST", encoded_path(request.target), request.headers, signed, request.body, _SECRET_KEY, scope
)
assert fields["Signature"] == expected[1], (request.target, fields)
def _assert_answered(response: httpx.Response, model: str, route: _Route = _SYSTEMONE) -> None:
assert response.status_code == 200, response.text
assert response.json() == route.answer, response.text
assert response.headers["x-litellm-model-group"] == model, dict(response.headers)
def _assert_runtime_target(request: Request, runtime: str, region: str = _US_EAST) -> None:
assert request.method == "POST", request.method
assert request.target == f"/runtimes/{quote(_arn(runtime, region), safe='')}/invocations", request.target
assert request.headers["host"] == _host(region), dict(request.headers)
def _assert_runtime_call(request: Request, runtime: str, region: str = _US_EAST, route: _Route = _SYSTEMONE) -> None:
_assert_runtime_target(request, runtime, region)
assert json.loads(request.body) == route.upstream, request.body
def _canonical(body: JsonValue | Mapping[str, JsonValue]) -> str:
return json.dumps(body, sort_keys=True)
def _spend_row(call_id: str) -> dict[str, JsonValue]:
rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70)
return rows[0]
def _spend_rows(call_ids: Sequence[str]) -> list[dict[str, JsonValue]]:
query: Final = (
"SELECT request_id, status FROM \"LiteLLM_SpendLogs\" WHERE request_id = ANY(string_to_array(%s, ','))"
)
return eventually(
lambda: read_rows(query, (",".join(call_ids),)), lambda found: len(found) >= len(call_ids), seconds=70
)
def _connect_targets(runtime: _Runtime) -> frozenset[str]:
return frozenset(runtime.tunnel.targets())
def test_an_arn_deployment_answers_on_every_decisions_route_signed_for_bedrock_agentcore(
agentcore: Gateway, runtime: _Runtime
) -> None:
runtime.reset()
responses: Final = tuple(_decide(agentcore, _SIGNED, route) for route in _ROUTES)
for response, route in zip(responses, _ROUTES, strict=True):
_assert_answered(response, _SIGNED, route)
assert _connect_targets(runtime) == {f"{_host(_US_EAST)}:443"}
calls: Final = runtime.calls()
assert len(calls) == len(_ROUTES), calls
for call, route in zip(calls, _ROUTES, strict=True):
_assert_runtime_call(call, "audit_signed", route=route)
assert call.headers[_SESSION_HEADER] == _default_session("audit_signed"), dict(call.headers)
_assert_signed(call)
for response in responses:
row = _spend_row(response.headers["x-litellm-call-id"])
assert row == {
"status": "success",
"call_type": "adecisions",
"custom_llm_provider": "strands_decider",
"model_group": _SIGNED,
"api_base": _invocations_url("audit_signed"),
"prompt_tokens": 211,
"completion_tokens": 2,
}, row
def test_the_arn_region_picks_the_host_and_signing_scope_over_aws_region_name(
agentcore: Gateway, runtime: _Runtime
) -> None:
runtime.reset()
_assert_answered(_decide(agentcore, _REGION), _REGION)
assert _connect_targets(runtime) == {f"{_host(_EU_WEST)}:443"}
(call,) = runtime.calls()
_assert_runtime_call(call, "audit_region", _EU_WEST)
_assert_signed(call, _EU_WEST)
def test_a_configured_session_id_is_sent_as_written_and_signed(agentcore: Gateway, runtime: _Runtime) -> None:
runtime.reset()
_assert_answered(_decide(agentcore, _SESSION), _SESSION)
_assert_answered(_decide(agentcore, _LONG), _LONG)
calls: Final = runtime.calls()
(configured,) = _calls_to(calls, "audit_session")
(long,) = _calls_to(calls, "audit_session_long")
assert configured.headers[_SESSION_HEADER] == _CONFIGURED_SESSION, dict(configured.headers)
assert long.headers[_SESSION_HEADER] == _LONG_SESSION
_assert_signed(configured)
_assert_signed(long)
def test_a_deployment_api_key_sends_bearer_and_the_session_without_sigv4(agentcore: Gateway, runtime: _Runtime) -> None:
runtime.reset()
_assert_answered(_decide(agentcore, _BEARER), _BEARER)
(call,) = runtime.calls()
_assert_runtime_call(call, "audit_bearer")
assert call.headers["authorization"] == f"Bearer {_JWT}", dict(call.headers)
assert call.headers[_SESSION_HEADER] == _default_session("audit_bearer"), dict(call.headers)
assert "x-amz-date" not in call.headers, dict(call.headers)
def test_reserved_extra_headers_are_dropped_and_the_rest_are_kept(agentcore: Gateway, runtime: _Runtime) -> None:
runtime.reset()
_assert_answered(_decide(agentcore, _HEADERS), _HEADERS)
caller_headers: Final[dict[str, JsonValue]] = {
"x-amz-security-token": "caller-forged-token",
"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": "caller-forged-session-0000000000000000001",
"x-audit-caller": "caller-trace",
}
_assert_answered(_decide(agentcore, _HEADERS, extra_headers=caller_headers), _HEADERS)
deployment, caller = runtime.calls()
for call in (deployment, caller):
_assert_runtime_call(call, "audit_headers")
assert call.headers[_SESSION_HEADER] == _default_session("audit_headers"), dict(call.headers)
assert "x-amz-target" not in call.headers, dict(call.headers)
assert "x-amz-security-token" not in call.headers, dict(call.headers)
_assert_signed(call)
assert deployment.headers["x-audit-trace"] == "deployment-trace", dict(deployment.headers)
assert caller.headers["x-audit-caller"] == "caller-trace", dict(caller.headers)
def test_an_opted_in_caller_session_id_becomes_the_header_and_stays_out_of_the_body(
agentcore: Gateway, runtime: _Runtime
) -> None:
runtime.reset()
chosen: Final = f"agentcore-audit-caller-session-{uuid.uuid4().hex}"
_assert_answered(_decide(agentcore, _OPT_IN, agentcore_runtime_session_id=chosen), _OPT_IN)
_assert_answered(_decide(agentcore, _OPT_IN), _OPT_IN)
first, second = runtime.calls()
for call in (first, second):
_assert_runtime_call(call, "audit_opt_in")
_assert_signed(call)
assert first.headers[_SESSION_HEADER] == chosen, dict(first.headers)
assert second.headers[_SESSION_HEADER] == _default_session("audit_opt_in"), dict(second.headers)
def test_an_arn_deployment_without_api_base_takes_the_runtime_from_the_environment(
agentcore: Gateway, runtime: _Runtime
) -> None:
runtime.reset()
response: Final = _decide(agentcore, _ENV)
_assert_answered(response, _ENV)
(call,) = runtime.calls()
_assert_runtime_call(call, "audit_env")
_assert_signed(call)
assert _spend_row(response.headers["x-litellm-call-id"])["api_base"] == _invocations_url("audit_env")
def test_an_arn_deployment_added_through_the_management_api_answers_signed(
agentcore: Gateway, runtime: _Runtime
) -> None:
runtime.reset()
with agentcore.scenario() as scenario:
model: Final = scenario.model(
model_info=None,
model=_MODEL,
api_base=_arn("audit_stored"),
api_key=None,
aws_access_key_id=_ACCESS_KEY,
aws_secret_access_key=_SECRET_KEY,
)
_assert_answered(_decide(agentcore, model), model)
(call,) = runtime.calls()
_assert_runtime_call(call, "audit_stored")
_assert_signed(call)
def test_health_and_test_connection_probe_the_runtime_signed(agentcore: Gateway, runtime: _Runtime) -> None:
runtime.reset()
health: Final = agentcore.request("GET", "/health", params={"model": _HEALTH})
assert health.status_code == 200, health.text
assert (health.json()["healthy_count"], health.json()["unhealthy_count"]) == (1, 0), health.text
(probe,) = runtime.calls()
assert probe.target == f"/runtimes/{quote(_arn('audit_health'), safe='')}/invocations", probe.target
_assert_signed(probe)
connection: Final = agentcore.request(
"POST",
"/health/test_connection",
{
"litellm_params": {"model": _MODEL, "api_base": _arn("audit_test_connection"), **_SIGNING},
"mode": "evaluation",
},
)
assert connection.status_code == 200, connection.text
assert connection.json()["status"] == "success", connection.text
(tested,) = runtime.calls()
assert _runtime_of(tested) == "audit_test_connection", tested.target
_assert_signed(tested)
@pytest.mark.parametrize(("code", "status"), tuple(_ERROR_STATUS.items()), ids=tuple(_ERROR_STATUS))
def test_a_runtime_error_reply_maps_to_its_status_with_the_runtime_message(
agentcore: Gateway, runtime: _Runtime, code: str, status: int
) -> None:
runtime.reset()
model: Final = f"agentcore-audit-error-{code}"
response: Final = _decide(agentcore, model)
assert response.status_code == status, response.text
assert f"Strands Decider runtime error '{code}': runtime refused with {code}" in response.text, response.text
assert len(_calls_to(runtime.calls(), f"audit_error_{code}")) == 1
row: Final = _spend_row(response.headers["x-litellm-call-id"])
assert (row["status"], row["model_group"]) == ("failure", model), row
@pytest.mark.parametrize("status", tuple(_HTTP_ERRORS), ids=tuple(str(status) for status in _HTTP_ERRORS))
def test_a_runtime_http_error_keeps_its_status_and_message(agentcore: Gateway, runtime: _Runtime, status: int) -> None:
runtime.reset()
model: Final = f"agentcore-audit-status-{status}"
response: Final = _decide(agentcore, model)
assert response.status_code == status, response.text
assert f"scripted {_HTTP_ERRORS[status]}" in response.text, response.text
assert len(_calls_to(runtime.calls(), f"audit_status_{status}")) == 1
assert _spend_row(response.headers["x-litellm-call-id"])["status"] == "failure"
@pytest.mark.parametrize("label", tuple(_MALFORMED_SESSIONS))
def test_a_malformed_configured_session_id_fails_only_its_own_deployment(
agentcore: Gateway, runtime: _Runtime, label: str
) -> None:
runtime.reset()
model: Final = f"agentcore-audit-session-{label}"
response: Final = _decide(agentcore, model)
calls: Final = _calls_to(runtime.calls(), f"audit_session_{label}")
if label == "empty":
_assert_answered(response, model)
(call,) = calls
assert call.headers[_SESSION_HEADER] == "", dict(call.headers)
else:
assert response.status_code == 500, response.text
assert "agentcore_runtime_session_id" in response.text, response.text
assert calls == ()
_assert_answered(_decide(agentcore, _SIGNED), _SIGNED)
def test_a_deployment_without_aws_credentials_fails_before_dialing_the_runtime(
agentcore: Gateway, runtime: _Runtime
) -> None:
runtime.reset()
response: Final = _decide(agentcore, _UNSIGNED)
assert response.status_code == 500, response.text
assert "credentials" in response.text.lower(), response.text
assert runtime.calls() == ()
assert _connect_targets(runtime) == frozenset()
_assert_answered(_decide(agentcore, _SIGNED), _SIGNED)
@pytest.mark.parametrize("label", tuple(_NEAR_MISS_ARNS))
def test_a_near_miss_arn_keeps_the_plain_api_base_path(agentcore: Gateway, runtime: _Runtime, label: str) -> None:
runtime.reset()
response: Final = _decide(agentcore, f"agentcore-audit-near-miss-{label}")
assert 500 <= response.status_code < 600, response.text
assert "litellm.APIConnectionError" in response.text, response.text
assert _connect_targets(runtime) == frozenset()
assert runtime.calls() == ()
def test_a_dropped_connection_is_a_connection_error_after_one_attempt(agentcore: Gateway, runtime: _Runtime) -> None:
runtime.reset()
response: Final = _decide(agentcore, _DROPPED)
assert response.status_code == 500, response.text
assert "litellm.APIConnectionError" in response.text, response.text
assert len(_calls_to(runtime.calls(), "audit_dropped")) == 1
assert _spend_row(response.headers["x-litellm-call-id"])["status"] == "failure"
def _burst_routes(size: int) -> tuple[_Route, ...]:
return tuple(_ROUTES[index % len(_ROUTES)] for index in range(size))
def _burst(gateway: Gateway, model: str, size: int) -> tuple[httpx.Response, ...]:
with ThreadPoolExecutor(max_workers=size) as pool:
return tuple(pool.map(partial(_decide, gateway, model), _burst_routes(size)))
def _health_counts(gateway: Gateway, model: str) -> tuple[int, int, int]:
response: Final = gateway.request("GET", "/health", params={"model": model})
body: Final = JSON_OBJECT.validate_json(response.content)
healthy: Final = body["healthy_count"]
unhealthy: Final = body["unhealthy_count"]
assert isinstance(healthy, int) and isinstance(unhealthy, int), body
return response.status_code, healthy, unhealthy
def test_a_runtime_outage_mid_burst_fails_each_call_once_and_recovers_signed(
agentcore: Gateway, runtime: _Runtime
) -> None:
runtime.reset()
runtime.outage.set()
failed: Final = _burst(agentcore, _OUTAGE, 10)
assert [response.status_code for response in failed] == [500] * 10, [response.text for response in failed]
assert _health_counts(agentcore, _OUTAGE) == (503, 0, 1)
runtime.outage.clear()
runtime.calls()
served: Final = _burst(agentcore, _OUTAGE, 20)
for response, route in zip(served, _burst_routes(20), strict=True):
_assert_answered(response, _OUTAGE, route)
calls: Final = runtime.calls()
assert len(calls) == 20, calls
for call in calls:
_assert_runtime_target(call, "audit_outage")
_assert_signed(call)
assert sorted(_canonical(_JSON.validate_json(call.body)) for call in calls) == sorted(
_canonical(route.upstream) for route in _burst_routes(20)
)
assert _health_counts(agentcore, _OUTAGE) == (200, 1, 0)
expected: Final = {
**{response.headers["x-litellm-call-id"]: "failure" for response in failed},
**{response.headers["x-litellm-call-id"]: "success" for response in served},
}
assert len(expected) == 30, expected
rows: Final = _spend_rows(tuple(expected))
assert sorted((str(row["request_id"]), row["status"]) for row in rows) == sorted(expected.items()), rows
async def test_sdk_sync_and_async_calls_reach_the_runtime_signed(
runtime: _Runtime, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
runtime.reset()
for name in _scrubbed():
monkeypatch.delenv(name)
for name, value in {
"HTTPS_PROXY": runtime.tunnel.url,
"NO_PROXY": "127.0.0.1,localhost",
"SSL_CERT_FILE": str(runtime.cert),
**_isolated_aws(tmp_path),
}.items():
monkeypatch.setenv(name, value)
litellm.in_memory_llm_clients_cache.flush_cache()
try:
synchronous: Final = litellm.decisions(
model=_MODEL,
state=_STATE,
questions=_SDK_QUESTIONS,
api_base=_arn("audit_sdk"),
aws_access_key_id=_ACCESS_KEY,
aws_secret_access_key=_SECRET_KEY,
)
asynchronous: Final = await litellm.adecisions(
model=_MODEL,
state=_STATE,
questions=_SDK_QUESTIONS,
api_base=_arn("audit_sdk"),
aws_access_key_id=_ACCESS_KEY,
aws_secret_access_key=_SECRET_KEY,
)
monkeypatch.setenv("STRANDS_DECIDER_API_KEY", _ENV_JWT)
bearer: Final = await litellm.adecisions(
model=_MODEL, state=_STATE, questions=_SDK_QUESTIONS, api_base=_arn("audit_sdk_env_key")
)
finally:
litellm.in_memory_llm_clients_cache.flush_cache()
for response in (synchronous, asynchronous, bearer):
assert response.model_dump(mode="json") == {"model": _BODY_MODEL, "answers": _ANSWERS, "usage": _USAGE}
calls: Final = runtime.calls()
signed: Final = _calls_to(calls, "audit_sdk")
assert len(signed) == 2, calls
for call in signed:
_assert_runtime_call(call, "audit_sdk")
assert call.headers[_SESSION_HEADER] == _default_session("audit_sdk"), dict(call.headers)
_assert_signed(call)
(env_key,) = _calls_to(calls, "audit_sdk_env_key")
assert env_key.headers["authorization"] == f"Bearer {_ENV_JWT}", dict(env_key.headers)
assert "x-amz-date" not in env_key.headers, dict(env_key.headers)
def _answering(request: Request) -> Reply:
return _json_reply({"model": _BODY_MODEL, "answers": _ANSWERS, "usage": _USAGE})
def _envelope(request: Request) -> Reply:
return _json_reply({"error": {"code": "bad_request", "message": "plain server refused"}})
def _plain_model(scenario: Scenario, wire: Wire, **extra: JsonValue) -> str:
return scenario.model(model_info=None, model=_MODEL, api_base=wire.url, api_key=None, **extra)
def test_a_request_body_session_id_is_refused_on_every_route_without_an_upstream_call(gateway: Gateway) -> None:
with wire_server(_answering) as wire, gateway.scenario() as scenario:
model: Final = _plain_model(scenario, wire)
responses: Final = tuple(
_decide(gateway, model, route, agentcore_runtime_session_id=f"caller-{uuid.uuid4().hex}")
for route in _ROUTES
)
for response in responses:
assert response.status_code == 401, response.text
assert "agentcore_runtime_session_id is not allowed in request body" in response.text, response.text
assert wire.drain() == ()
def test_a_plain_api_base_opted_into_the_session_param_sends_no_session_header_or_body_field(
gateway: Gateway,
) -> None:
with wire_server(_answering) as wire, gateway.scenario() as scenario:
model: Final = _plain_model(
scenario, wire, configurable_clientside_auth_params=["agentcore_runtime_session_id"]
)
_assert_answered(_decide(gateway, model, agentcore_runtime_session_id=f"caller-{uuid.uuid4().hex}"), model)
(call,) = wire.drain()
assert call.target == "/v1/systemone", call.target
assert _SESSION_HEADER not in call.headers, dict(call.headers)
assert json.loads(call.body) == _UPSTREAM_BODY, call.body
def test_a_plain_api_base_error_envelope_keeps_the_unexpected_response_server_error(gateway: Gateway) -> None:
with wire_server(_envelope) as wire, gateway.scenario() as scenario:
model: Final = _plain_model(scenario, wire)
response: Final = _decide(gateway, model)
assert response.status_code == 500, response.text
assert "Strands Decider runtime error" not in response.text, response.text
assert len(wire.drain()) == 1
def test_a_plain_api_base_forwards_amazon_shaped_extra_headers_unchanged(gateway: Gateway) -> None:
headers: Final[dict[str, JsonValue]] = {
"x-amz-target": "plain-target",
"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": "plain-session-kept",
"x-audit-trace": "plain-trace",
}
with wire_server(_answering) as wire, gateway.scenario() as scenario:
model: Final = _plain_model(scenario, wire, extra_headers=headers)
_assert_answered(_decide(gateway, model), model)
(call,) = wire.drain()
assert (call.headers["x-amz-target"], call.headers[_SESSION_HEADER], call.headers["x-audit-trace"]) == (
"plain-target",
"plain-session-kept",
"plain-trace",
), dict(call.headers)
assert "authorization" not in call.headers, dict(call.headers)
assert json.loads(call.body) == _UPSTREAM_BODY, call.body
def _worst_case_api_base() -> str:
return "arn:aws" + "-a" * 131072 + "!"
def test_a_worst_case_runtime_arn_pattern_input_is_classified_in_linear_time(
gateway: Gateway, record_property: Callable[[str, object], None]
) -> None:
worst: Final = _worst_case_api_base()
started: Final = time.perf_counter()
url: Final = StrandsDeciderDecisionsConfig().get_complete_url(worst, _BODY_MODEL)
classified: Final = time.perf_counter() - started
record_property("classifier_seconds", f"{classified:.6f}")
assert url == f"{worst}/v1/systemone"
assert classified < 0.5, classified
with wire_server(_answering) as wire, gateway.scenario() as scenario:
model: Final = _plain_model(scenario, wire, configurable_clientside_auth_params=["api_base"])
with ThreadPoolExecutor(max_workers=2) as pool:
timed: Final = pool.submit(_timed_decide, gateway, model, worst)
liveliness_started: Final = time.perf_counter()
liveliness: Final = gateway.request("GET", "/health/liveliness")
liveliness_seconds: Final = time.perf_counter() - liveliness_started
response, decide_seconds = timed.result()
record_property("decide_seconds", f"{decide_seconds:.6f}")
record_property("liveliness_seconds", f"{liveliness_seconds:.6f}")
assert liveliness.status_code == 200, liveliness.text
assert 500 <= response.status_code < 600, response.text
assert "litellm.APIConnectionError" in response.text, response.text[:2000]
assert decide_seconds < 5, decide_seconds
assert liveliness_seconds < 2, liveliness_seconds
assert wire.drain() == ()
def _timed_decide(gateway: Gateway, model: str, api_base: str) -> tuple[httpx.Response, float]:
started: Final = time.perf_counter()
response: Final = _decide(gateway, model, api_base=api_base)
return response, time.perf_counter() - started

647
tests/integration/routing/test_per_model_rate_limit_fallbacks.py Normal file
View file

@ -0,0 +1,647 @@
from __future__ import annotations
import json
import re
import signal
import uuid
from collections import Counter
from collections.abc import Iterator, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from contextlib import ExitStack
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
import anthropic
import httpx
import openai
import psutil
import pytest
import yaml
from integration._support.client import (
JSON_OBJECT,
Gateway,
Scenario,
eventually,
gateway_from_environment,
object_value,
string_value,
)
from integration._support.openai_wire import chat_reply, responses_reply
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from integration.routing.test_payment_required_mapping import OPENAI_MODEL, is_model_info_probe, model_list_reply
from pydantic import JsonValue, TypeAdapter
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
SETTING: Final = "disable_fallbacks_on_per_model_rate_limits"
FRESH_CONNECTION: Final = MappingProxyType({"Connection": "close"})
MARKER: Final = re.compile(rb"pmrl-[0-9a-f]{32}")
WORKER_STARTED: Final = re.compile(r"Started server process \[(\d+)\]")
STARTUP_COMPLETE: Final = "Application startup complete."
CONFIG_LIST: Final = TypeAdapter(list[dict[str, JsonValue]])
WORKER_PIDS: Final = TypeAdapter(tuple[int, ...])
BURST: Final = 8
TOKEN_LIMIT: Final = 100
INPUT_TOKEN_LIMIT: Final = 120
MAX_TOKENS: Final = 60
METERED_USAGE: Final = MappingProxyType({"prompt_tokens": 900, "completion_tokens": 100, "total_tokens": 1000})
MALFORMED_VALUES: Final[tuple[tuple[str, JsonValue], ...]] = (
("word", "sometimes"),
("empty", ""),
("five_kb", "x" * 5120),
("list", [True]),
("dict", {"enabled": True}),
("int", 2),
)
@dataclass(frozen=True, slots=True)
class Groups:
primary: str
fallback: str
last: str
metered: str
gated_hard: str
gated_soft: str
capped: str
def names(self) -> tuple[str, ...]:
return (self.primary, self.fallback, self.last, self.metered, self.gated_hard, self.gated_soft, self.capped)
@dataclass(frozen=True, slots=True)
class Rig:
gateway: Gateway
log: Path
groups: Groups
wires: Mapping[str, Wire]
def traffic(self) -> dict[str, tuple[str, ...]]:
return {group: markers_on(wire) for group, wire in self.wires.items()}
def expected(self, served: Mapping[str, Sequence[str]]) -> dict[str, tuple[str, ...]]:
return {group: tuple(sorted(served.get(group, ()))) for group in self.wires}
def new_marker() -> str:
return f"pmrl-{uuid.uuid4().hex}"
def new_groups() -> Groups:
suffix: Final = uuid.uuid4().hex[:10]
return Groups(
primary=f"per-model-primary-{suffix}",
fallback=f"per-model-fallback-{suffix}",
last=f"per-model-last-{suffix}",
metered=f"per-model-metered-{suffix}",
gated_hard=f"per-model-gated-hard-{suffix}",
gated_soft=f"per-model-gated-soft-{suffix}",
capped=f"per-model-capped-{suffix}",
)
def marker_in(request: Request) -> str:
found: Final = MARKER.search(request.body)
assert found is not None, (request.method, request.target, request.body[:300])
return found.group(0).decode()
def markers_on(wire: Wire) -> tuple[str, ...]:
return tuple(sorted(marker_in(request) for request in wire.drain() if not is_model_info_probe(request)))
def streamed(request: Request) -> bool:
return JSON_OBJECT.validate_json(request.body).get("stream") is True
def healthy_reply(request: Request) -> Reply:
if is_model_info_probe(request):
return model_list_reply()
marker: Final = marker_in(request)
if request.target.endswith("/responses"):
return responses_reply(f"resp_{marker}", OPENAI_MODEL, f"served {marker}", stream=streamed(request))
return chat_reply(f"chatcmpl-{marker}", OPENAI_MODEL, f"served {marker}", stream=streamed(request))
def metered_reply(request: Request) -> Reply:
if is_model_info_probe(request):
return model_list_reply()
marker: Final = marker_in(request)
body: Final = {
"id": f"chatcmpl-{marker}",
"object": "chat.completion",
"created": 1,
"model": OPENAI_MODEL,
"choices": [
{"index": 0, "message": {"role": "assistant", "content": f"served {marker}"}, "finish_reason": "stop"}
],
"usage": dict(METERED_USAGE),
}
return Reply(body=json.dumps(body).encode())
def model_entry(group: str, wire: Wire, **extra: JsonValue) -> dict[str, JsonValue]:
return {
"model_name": group,
"litellm_params": {
"model": OPENAI_MODEL,
"api_base": wire.url + "/v1",
"api_key": "synthetic-openai-key",
**extra,
},
}
def proxy_config(
directory: Path,
*,
model_list: Sequence[Mapping[str, JsonValue]],
fallbacks: Sequence[Mapping[str, JsonValue]],
general_settings: Mapping[str, JsonValue],
callbacks: Sequence[str] = (),
) -> Path:
base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
litellm_settings: Final = object_value(base["litellm_settings"])
config: Final = {
**base,
"model_list": [dict(entry) for entry in model_list],
"general_settings": {**object_value(base["general_settings"]), **general_settings},
"litellm_settings": {**litellm_settings, **({"callbacks": list(callbacks)} if callbacks else {})},
"router_settings": {
**object_value(base["router_settings"]),
"num_retries": 0,
"fallbacks": [dict(entry) for entry in fallbacks],
},
}
path: Final = directory / f"per-model-limits-{uuid.uuid4().hex[:8]}.yaml"
path.write_text(yaml.safe_dump(config))
return path
def enforcing_config(directory: Path, groups: Groups, wires: Mapping[str, Wire]) -> Path:
return proxy_config(
directory,
model_list=(
model_entry(groups.primary, wires[groups.primary]),
model_entry(groups.fallback, wires[groups.fallback]),
model_entry(groups.last, wires[groups.last]),
model_entry(groups.metered, wires[groups.metered]),
model_entry(groups.gated_hard, wires[groups.gated_hard], rpm=1),
model_entry(groups.gated_soft, wires[groups.gated_soft], rpm=1),
model_entry(groups.capped, wires[groups.capped]),
),
fallbacks=(
{groups.primary: [groups.fallback]},
{groups.metered: [groups.fallback]},
{groups.gated_hard: [groups.capped, groups.last]},
{groups.gated_soft: [groups.fallback]},
),
general_settings={SETTING: True},
callbacks=("dynamic_rate_limiter_v3",),
)
@pytest.fixture(scope="module")
def enforcing(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
directory: Final = tmp_path_factory.mktemp("per-model-limits")
groups: Final = new_groups()
with ExitStack() as stack:
rig_gateway: Final = stack.enter_context(gateway_from_environment())
wires: Final = MappingProxyType(
{
group: stack.enter_context(wire_server(metered_reply if group == groups.metered else healthy_reply))
for group in groups.names()
}
)
owned: Final = stack.enter_context(
owned_proxy_process(
rig_gateway, directory, {}, config=enforcing_config(directory, groups, wires), workers=2
)
)
yield Rig(owned.gateway, owned.log, groups, wires)
def chat(gateway: Gateway, key: str, group: str, content: str, **extra: JsonValue) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{"model": group, "messages": [{"role": "user", "content": content}], **extra},
key=key,
headers=FRESH_CONNECTION,
)
def responses_call(gateway: Gateway, key: str, group: str, marker: str, *, stream: bool) -> httpx.Response:
return gateway.request(
"POST", "/v1/responses", {"model": group, "input": marker, "stream": stream}, key=key, headers=FRESH_CONNECTION
)
def openai_sdk(gateway: Gateway, key: str) -> openai.OpenAI:
return openai.OpenAI(
base_url=str(gateway.client.base_url).rstrip("/") + "/v1",
api_key=key,
max_retries=0,
default_headers=dict(FRESH_CONNECTION),
)
def openai_async_sdk(gateway: Gateway, key: str) -> openai.AsyncOpenAI:
return openai.AsyncOpenAI(
base_url=str(gateway.client.base_url).rstrip("/") + "/v1",
api_key=key,
max_retries=0,
default_headers=dict(FRESH_CONNECTION),
)
def anthropic_sdk(gateway: Gateway, key: str) -> anthropic.Anthropic:
return anthropic.Anthropic(
base_url=str(gateway.client.base_url).rstrip("/"),
api_key=key,
max_retries=0,
default_headers=dict(FRESH_CONNECTION),
)
def anthropic_async_sdk(gateway: Gateway, key: str) -> anthropic.AsyncAnthropic:
return anthropic.AsyncAnthropic(
base_url=str(gateway.client.base_url).rstrip("/"),
api_key=key,
max_retries=0,
default_headers=dict(FRESH_CONNECTION),
)
def error_message(body: str) -> str:
return string_value(object_value(JSON_OBJECT.validate_json(body)["error"])["message"])
def assert_limit_message(body: str, descriptor: str, group: str, kind: str, limit: int) -> None:
message: Final = error_message(body)
assert f"Rate limit exceeded for {descriptor}: " in message, message
assert f":{group}. Limit type: {kind}. Current limit: {limit}, Remaining: " in message, message
def assert_refused(
response: httpx.Response, descriptor: str, group: str, *, kind: str = "requests", limit: int = 1
) -> None:
assert response.status_code == 429, (
response.status_code,
response.headers.get("x-litellm-model-group"),
response.text,
)
assert_limit_message(response.text, descriptor, group, kind, limit)
def assert_served(response: httpx.Response, group: str, identity: str) -> None:
assert response.status_code == 200, response.text
assert response.headers["x-litellm-model-group"] == group, (
response.headers.get("x-litellm-model-group"),
response.text,
)
assert identity in response.text, response.text
def burst(gateway: Gateway, calls: Sequence[tuple[str, str]], group: str) -> tuple[httpx.Response, ...]:
def send(call: tuple[str, str]) -> httpx.Response:
return chat(gateway, call[0], group, call[1])
with ThreadPoolExecutor(max_workers=len(calls)) as pool:
return tuple(pool.map(send, calls))
def served_once_in_a_burst(rig: Rig, markers: Sequence[str], responses: Sequence[httpx.Response]) -> str:
assert Counter(response.status_code for response in responses) == Counter({200: 1, 429: BURST - 1}), tuple(
(response.status_code, response.headers.get("x-litellm-model-group"), response.text[:200])
for response in responses
)
refusals: Final = tuple(response for response in responses if response.status_code == 429)
for response in refusals:
assert_refused(response, "model_per_key", rig.groups.primary)
(served,) = (marker for marker, response in zip(markers, responses) if response.status_code == 200)
return served
def worker_startups(log: Path) -> tuple[tuple[int, ...], int]:
text: Final = log.read_text(errors="replace")
return WORKER_PIDS.validate_python(WORKER_STARTED.findall(text)), text.count(STARTUP_COMPLETE)
def listed_setting(gateway: Gateway) -> tuple[dict[str, JsonValue], ...]:
response: Final = gateway.request("GET", "/config/list", params={"config_type": "general_settings"})
assert response.status_code == 200, response.text
return tuple(entry for entry in CONFIG_LIST.validate_json(response.content) if entry["field_name"] == SETTING)
def update_setting(gateway: Gateway, value: JsonValue) -> httpx.Response:
return gateway.request(
"POST",
"/config/field/update",
{"field_name": SETTING, "field_value": value, "config_type": "general_settings"},
)
def scope_key_fields(scenario: Scenario, scope: str, group: str) -> dict[str, JsonValue]:
limit: Final[dict[str, JsonValue]] = {group: 1}
if scope == "team":
return {"team_id": scenario.team(model_rpm_limit=limit)}
if scope == "organization":
return {"team_id": scenario.team(organization_id=scenario.organization(model_rpm_limit=limit))}
team: Final = scenario.team()
return {"team_id": team, "project_id": scenario.project(team, model_rpm_limit=limit)}
def test_chat_openai_sdk_over_key_model_rpm_gets_429_instead_of_the_fallback(enforcing: Rig) -> None:
enforcing.traffic()
primary: Final = enforcing.groups.primary
first, second = new_marker(), new_marker()
with enforcing.gateway.scenario() as scenario:
client: Final = openai_sdk(enforcing.gateway, scenario.key(model_rpm_limit={primary: 1}))
served: Final = client.chat.completions.with_raw_response.create(
model=primary, messages=[{"role": "user", "content": first}]
)
assert served.headers["x-litellm-model-group"] == primary
assert served.parse().id == f"chatcmpl-{first}"
with pytest.raises(openai.RateLimitError) as refused:
client.chat.completions.create(model=primary, messages=[{"role": "user", "content": second}])
assert_limit_message(refused.value.response.text, "model_per_key", primary, "requests", 1)
assert enforcing.traffic() == enforcing.expected({primary: (first,)})
async def test_chat_stream_openai_async_sdk_over_key_model_rpm_gets_429_instead_of_the_fallback(
enforcing: Rig,
) -> None:
enforcing.traffic()
primary: Final = enforcing.groups.primary
first, second = new_marker(), new_marker()
with enforcing.gateway.scenario() as scenario:
async with openai_async_sdk(enforcing.gateway, scenario.key(model_rpm_limit={primary: 1})) as client:
stream: Final = await client.chat.completions.create(
model=primary, messages=[{"role": "user", "content": first}], stream=True
)
identities: Final = {chunk.id async for chunk in stream}
assert identities == {f"chatcmpl-{first}"}
with pytest.raises(openai.RateLimitError) as refused:
await client.chat.completions.create(
model=primary, messages=[{"role": "user", "content": second}], stream=True
)
assert_limit_message(refused.value.response.text, "model_per_key", primary, "requests", 1)
assert enforcing.traffic() == enforcing.expected({primary: (first,)})
def test_messages_anthropic_sdk_over_key_model_rpm_gets_429_instead_of_the_fallback(enforcing: Rig) -> None:
enforcing.traffic()
primary: Final = enforcing.groups.primary
first, second = new_marker(), new_marker()
with enforcing.gateway.scenario() as scenario:
client: Final = anthropic_sdk(enforcing.gateway, scenario.key(model_rpm_limit={primary: 1}))
served: Final = client.messages.with_raw_response.create(
model=primary, max_tokens=32, messages=[{"role": "user", "content": first}]
)
assert served.headers["x-litellm-model-group"] == primary
block: Final = served.parse().content[0]
assert isinstance(block, anthropic.types.TextBlock), block
assert block.text == f"served {first}"
with pytest.raises(anthropic.RateLimitError) as refused:
client.messages.create(model=primary, max_tokens=32, messages=[{"role": "user", "content": second}])
assert_limit_message(refused.value.response.text, "model_per_key", primary, "requests", 1)
assert enforcing.traffic() == enforcing.expected({primary: (first,)})
async def test_messages_stream_anthropic_async_sdk_over_key_model_rpm_gets_429_instead_of_the_fallback(
enforcing: Rig,
) -> None:
enforcing.traffic()
primary: Final = enforcing.groups.primary
first, second = new_marker(), new_marker()
with enforcing.gateway.scenario() as scenario:
async with anthropic_async_sdk(enforcing.gateway, scenario.key(model_rpm_limit={primary: 1})) as client:
async with client.messages.stream(
model=primary, max_tokens=32, messages=[{"role": "user", "content": first}]
) as served:
text: Final = await served.get_final_text()
assert text == f"served {first}"
with pytest.raises(anthropic.RateLimitError) as refused:
async with client.messages.stream(
model=primary, max_tokens=32, messages=[{"role": "user", "content": second}]
) as rejected:
await rejected.get_final_text()
assert_limit_message(refused.value.response.text, "model_per_key", primary, "requests", 1)
assert enforcing.traffic() == enforcing.expected({primary: (first,)})
@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream"))
def test_responses_httpx_over_key_model_rpm_gets_429_instead_of_the_fallback(enforcing: Rig, stream: bool) -> None:
enforcing.traffic()
primary: Final = enforcing.groups.primary
first, second = new_marker(), new_marker()
with enforcing.gateway.scenario() as scenario:
key: Final = scenario.key(model_rpm_limit={primary: 1})
assert_served(responses_call(enforcing.gateway, key, primary, first, stream=stream), primary, f"resp_{first}")
assert_refused(responses_call(enforcing.gateway, key, primary, second, stream=stream), "model_per_key", primary)
assert enforcing.traffic() == enforcing.expected({primary: (first,)})
@pytest.mark.parametrize(
("scope", "descriptor"),
(("team", "model_per_team"), ("organization", "model_per_organization"), ("project", "model_per_project")),
ids=("team", "organization", "project"),
)
def test_chat_over_a_scope_model_rpm_gets_429_instead_of_the_fallback(
enforcing: Rig, scope: str, descriptor: str
) -> None:
enforcing.traffic()
primary: Final = enforcing.groups.primary
first, second = new_marker(), new_marker()
with enforcing.gateway.scenario() as scenario:
fields: Final = scope_key_fields(scenario, scope, primary)
first_key, second_key = scenario.key(**fields), scenario.key(**fields)
assert_served(chat(enforcing.gateway, first_key, primary, first), primary, f"chatcmpl-{first}")
assert_refused(chat(enforcing.gateway, second_key, primary, second), descriptor, primary)
assert enforcing.traffic() == enforcing.expected({primary: (first,)})
def test_chat_over_project_model_itpm_gets_429_instead_of_the_fallback(enforcing: Rig) -> None:
enforcing.traffic()
metered: Final = enforcing.groups.metered
first, second = new_marker(), new_marker()
filler: Final = " word" * 50
with enforcing.gateway.scenario() as scenario:
team: Final = scenario.team()
key: Final = scenario.key(
team_id=team, project_id=scenario.project(team, model_itpm_limit={metered: INPUT_TOKEN_LIMIT})
)
assert_served(chat(enforcing.gateway, key, metered, first + filler), metered, f"chatcmpl-{first}")
assert_refused(
chat(enforcing.gateway, key, metered, second + filler),
"model_per_project_itpm",
metered,
kind="tokens",
limit=INPUT_TOKEN_LIMIT,
)
assert enforcing.traffic() == enforcing.expected({metered: (first,)})
def test_chat_over_key_model_tpm_gets_429_instead_of_the_fallback(enforcing: Rig) -> None:
enforcing.traffic()
metered: Final = enforcing.groups.metered
first, second = new_marker(), new_marker()
with enforcing.gateway.scenario() as scenario:
key: Final = scenario.key(model_tpm_limit={metered: TOKEN_LIMIT})
assert_served(chat(enforcing.gateway, key, metered, first, max_tokens=MAX_TOKENS), metered, f"chatcmpl-{first}")
assert_refused(
chat(enforcing.gateway, key, metered, second, max_tokens=MAX_TOKENS),
"model_per_key",
metered,
kind="tokens",
limit=TOKEN_LIMIT,
)
assert enforcing.traffic() == enforcing.expected({metered: (first,)})
def test_key_router_settings_fallbacks_are_skipped_on_a_per_model_429(enforcing: Rig) -> None:
enforcing.traffic()
primary: Final = enforcing.groups.primary
first, second = new_marker(), new_marker()
with enforcing.gateway.scenario() as scenario:
key: Final = scenario.key(
model_rpm_limit={primary: 1}, router_settings={"fallbacks": [{primary: [enforcing.groups.last]}]}
)
assert_served(chat(enforcing.gateway, key, primary, first), primary, f"chatcmpl-{first}")
assert_refused(chat(enforcing.gateway, key, primary, second), "model_per_key", primary)
assert enforcing.traffic() == enforcing.expected({primary: (first,)})
def test_capacity_429_falls_back_but_a_capped_fallback_answers_429_instead_of_the_next_one(enforcing: Rig) -> None:
enforcing.traffic()
groups: Final = enforcing.groups
warm, spend, refused = new_marker(), new_marker(), new_marker()
with enforcing.gateway.scenario() as scenario:
key: Final = scenario.key(model_rpm_limit={groups.capped: 1})
assert_served(chat(enforcing.gateway, key, groups.gated_hard, warm), groups.gated_hard, f"chatcmpl-{warm}")
assert_served(chat(enforcing.gateway, key, groups.capped, spend), groups.capped, f"chatcmpl-{spend}")
assert_refused(chat(enforcing.gateway, key, groups.gated_hard, refused), "model_per_key", groups.capped)
assert enforcing.traffic() == enforcing.expected({groups.gated_hard: (warm,), groups.capped: (spend,)})
def test_model_capacity_429_without_a_per_model_limit_still_falls_back(enforcing: Rig) -> None:
enforcing.traffic()
groups: Final = enforcing.groups
warm, moved = new_marker(), new_marker()
with enforcing.gateway.scenario() as scenario:
key: Final = scenario.key()
assert_served(chat(enforcing.gateway, key, groups.gated_soft, warm), groups.gated_soft, f"chatcmpl-{warm}")
assert_served(chat(enforcing.gateway, key, groups.gated_soft, moved), groups.fallback, f"chatcmpl-{moved}")
assert enforcing.traffic() == enforcing.expected({groups.gated_soft: (warm,), groups.fallback: (moved,)})
def test_key_rpm_limit_without_a_model_scope_answers_429_on_both_the_primary_and_the_fallback(enforcing: Rig) -> None:
enforcing.traffic()
primary: Final = enforcing.groups.primary
first, second = new_marker(), new_marker()
with enforcing.gateway.scenario() as scenario:
key: Final = scenario.key(rpm_limit=1)
assert_served(chat(enforcing.gateway, key, primary, first), primary, f"chatcmpl-{first}")
refused: Final = chat(enforcing.gateway, key, primary, second)
assert refused.status_code == 429, refused.text
assert error_message(refused.text).startswith("Rate limit exceeded for api_key: "), refused.text
assert ". Limit type: requests. Current limit: 1, Remaining: 0." in refused.text, refused.text
assert enforcing.traffic() == enforcing.expected({primary: (first,)})
def test_response_cache_hits_count_toward_key_model_rpm_and_the_next_request_gets_429(enforcing: Rig) -> None:
enforcing.traffic()
primary: Final = enforcing.groups.primary
marker: Final = new_marker()
with enforcing.gateway.scenario() as scenario:
key: Final = scenario.key(model_rpm_limit={primary: 2})
assert_served(chat(enforcing.gateway, key, primary, marker), primary, f"chatcmpl-{marker}")
assert_served(chat(enforcing.gateway, key, primary, marker), primary, f"chatcmpl-{marker}")
assert_refused(chat(enforcing.gateway, key, primary, marker), "model_per_key", primary, limit=2)
assert enforcing.traffic() == enforcing.expected({primary: (marker,)})
def test_concurrent_burst_over_key_model_rpm_serves_one_and_refuses_the_rest_without_fallback(enforcing: Rig) -> None:
enforcing.traffic()
primary: Final = enforcing.groups.primary
markers: Final = tuple(new_marker() for _ in range(BURST))
bystander: Final = new_marker()
with enforcing.gateway.scenario() as scenario:
limited: Final = scenario.key(model_rpm_limit={primary: 1})
unlimited: Final = scenario.key()
calls: Final = (*((limited, marker) for marker in markers), (unlimited, bystander))
responses: Final = burst(enforcing.gateway, calls, primary)
served: Final = served_once_in_a_burst(enforcing, markers, responses[:BURST])
assert_served(responses[BURST], primary, f"chatcmpl-{bystander}")
assert enforcing.traffic() == enforcing.expected({primary: (served, bystander)})
def test_yaml_owned_setting_refuses_a_database_write_and_keeps_enforcing(enforcing: Rig) -> None:
enforcing.traffic()
primary: Final = enforcing.groups.primary
refused_write: Final = update_setting(enforcing.gateway, False)
assert refused_write.status_code == 400, refused_write.text
detail: Final = object_value(JSON_OBJECT.validate_json(refused_write.content)["detail"])
assert detail["error"] == (
f"general_settings key '{SETTING}' is set in the config file and cannot be changed here."
), detail
assert detail["keys"] == [SETTING], detail
listed: Final = listed_setting(enforcing.gateway)
assert [(entry["field_type"], entry["field_value"], entry["editable"]) for entry in listed] == [
("Boolean", True, False)
], listed
first, second = new_marker(), new_marker()
with enforcing.gateway.scenario() as scenario:
key: Final = scenario.key(model_rpm_limit={primary: 1})
assert_served(chat(enforcing.gateway, key, primary, first), primary, f"chatcmpl-{first}")
assert_refused(chat(enforcing.gateway, key, primary, second), "model_per_key", primary)
assert enforcing.traffic() == enforcing.expected({primary: (first,)})
def test_worker_sigkill_and_replacement_keeps_refusing_per_model_overflow_without_fallback(enforcing: Rig) -> None:
enforcing.traffic()
pids, ready = eventually(
lambda: worker_startups(enforcing.log), lambda found: len(found[0]) >= 2 and found[1] >= 2, seconds=60
)
victim: Final = next(pid for pid in pids if psutil.pid_exists(pid))
psutil.Process(victim).send_signal(signal.SIGKILL)
eventually(
lambda: worker_startups(enforcing.log),
lambda found: len(found[0]) > len(pids) and found[1] > ready,
seconds=graceful_stop_seconds(),
)
primary: Final = enforcing.groups.primary
markers: Final = tuple(new_marker() for _ in range(BURST))
with enforcing.gateway.scenario() as scenario:
limited: Final = scenario.key(model_rpm_limit={primary: 1})
responses: Final = burst(enforcing.gateway, tuple((limited, marker) for marker in markers), primary)
served: Final = served_once_in_a_burst(enforcing, markers, responses)
assert enforcing.traffic() == enforcing.expected({primary: (served,)})
def test_config_list_reports_the_setting_as_an_unset_boolean(gateway: Gateway) -> None:
listed: Final = listed_setting(gateway)
assert [(entry["field_type"], entry["field_value"], entry["stored_in_db"]) for entry in listed] == [
("Boolean", None, None)
], listed
@pytest.mark.parametrize(
"value", tuple(value for _, value in MALFORMED_VALUES), ids=tuple(name for name, _ in MALFORMED_VALUES)
)
def test_config_field_update_rejects_a_non_boolean_value_and_stores_nothing(gateway: Gateway, value: JsonValue) -> None:
response: Final = update_setting(gateway, value)
try:
assert response.status_code == 400, response.text
assert JSON_OBJECT.validate_json(response.content) == {
"detail": {"error": f"Invalid type of field value={type(value)} passed in."}
}, response.text
listed: Final = listed_setting(gateway)
assert [(entry["field_value"], entry["stored_in_db"]) for entry in listed] == [(None, None)], listed
finally:
if response.status_code == 200:
gateway.post("/config/field/delete", {"field_name": SETTING, "config_type": "general_settings"})

194
tests/integration/routing/test_per_model_rate_limit_fallbacks_lifecycle.py Normal file
View file

@ -0,0 +1,194 @@
from __future__ import annotations
from collections.abc import Generator, Iterator, Mapping, Sequence
from contextlib import ExitStack, contextmanager
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
import pytest
from integration._support.client import eventually, gateway_from_environment, string_value
from integration._support.database import scratch_database
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Wire, wire_server
from integration.routing.test_per_model_rate_limit_fallbacks import (
SETTING,
Groups,
Rig,
assert_refused,
assert_served,
chat,
error_message,
healthy_reply,
listed_setting,
model_entry,
new_groups,
new_marker,
proxy_config,
update_setting,
)
from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(4 * graceful_stop_seconds() + 120)
NO_OVERRIDES: Final[Mapping[str, str]] = MappingProxyType({})
LEGACY_LIMITER: Final = MappingProxyType({"LEGACY_MULTI_INSTANCE_RATE_LIMITING": "true"})
LEGACY_ATTEMPTS: Final = 10
NOT_A_BOOLEAN: Final = "is not a boolean, treating it as disabled"
def serving_groups(groups: Groups) -> tuple[str, ...]:
return (groups.primary, groups.fallback, groups.last)
def fallback_config(
directory: Path, groups: Groups, wires: Mapping[str, Wire], general_settings: Mapping[str, JsonValue]
) -> Path:
return proxy_config(
directory,
model_list=tuple(model_entry(group, wires[group]) for group in serving_groups(groups)),
fallbacks=({groups.primary: [groups.fallback]},),
general_settings=general_settings,
)
@contextmanager
def fallback_rig(
directory: Path, general_settings: Mapping[str, JsonValue], environment: Mapping[str, str] = NO_OVERRIDES
) -> Generator[Rig]:
groups: Final = new_groups()
with ExitStack() as stack:
rig_gateway: Final = stack.enter_context(gateway_from_environment())
wires: Final = MappingProxyType(
{group: stack.enter_context(wire_server(healthy_reply)) for group in serving_groups(groups)}
)
config: Final = fallback_config(directory, groups, wires, general_settings)
owned: Final = stack.enter_context(
owned_proxy_process(rig_gateway, directory, environment, config=config, workers=2)
)
yield Rig(owned.gateway, owned.log, groups, wires)
@pytest.fixture(scope="module")
def lenient(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
with fallback_rig(tmp_path_factory.mktemp("per-model-limits-lenient"), {SETTING: None}) as rig:
yield rig
def through_first_refusal(responses: Iterator[httpx.Response]) -> Iterator[httpx.Response]:
for response in responses:
yield response
if response.status_code != 200:
return
def sent_until_refused(rig: Rig, key: str, markers: Sequence[str]) -> tuple[httpx.Response, ...]:
return tuple(through_first_refusal(chat(rig.gateway, key, rig.groups.primary, marker) for marker in markers))
def test_null_setting_keeps_falling_back_on_a_per_model_429_without_a_warning(lenient: Rig) -> None:
lenient.traffic()
primary, fallback = lenient.groups.primary, lenient.groups.fallback
first, second = new_marker(), new_marker()
with lenient.gateway.scenario() as scenario:
key: Final = scenario.key(model_rpm_limit={primary: 1})
assert_served(chat(lenient.gateway, key, primary, first), primary, f"chatcmpl-{first}")
assert_served(chat(lenient.gateway, key, primary, second), fallback, f"chatcmpl-{second}")
assert lenient.traffic() == lenient.expected({primary: (first,), fallback: (second,)})
assert NOT_A_BOOLEAN not in lenient.log.read_text(errors="replace")
def test_key_router_settings_fallbacks_still_serve_a_per_model_429_when_the_setting_is_off(lenient: Rig) -> None:
lenient.traffic()
primary, last = lenient.groups.primary, lenient.groups.last
first, second = new_marker(), new_marker()
with lenient.gateway.scenario() as scenario:
key: Final = scenario.key(model_rpm_limit={primary: 1}, router_settings={"fallbacks": [{primary: [last]}]})
assert_served(chat(lenient.gateway, key, primary, first), primary, f"chatcmpl-{first}")
assert_served(chat(lenient.gateway, key, primary, second), last, f"chatcmpl-{second}")
assert lenient.traffic() == lenient.expected({primary: (first,), last: (second,)})
def test_request_disable_fallbacks_answers_a_per_model_429_when_the_setting_is_off(lenient: Rig) -> None:
lenient.traffic()
primary: Final = lenient.groups.primary
first, second = new_marker(), new_marker()
with lenient.gateway.scenario() as scenario:
key: Final = scenario.key(model_rpm_limit={primary: 1})
assert_served(chat(lenient.gateway, key, primary, first), primary, f"chatcmpl-{first}")
assert_refused(chat(lenient.gateway, key, primary, second, disable_fallbacks=True), "model_per_key", primary)
assert lenient.traffic() == lenient.expected({primary: (first,)})
def test_yaml_true_string_turns_the_setting_on(tmp_path: Path) -> None:
with fallback_rig(tmp_path, {SETTING: "true"}) as rig, rig.gateway.scenario() as scenario:
primary: Final = rig.groups.primary
first, second = new_marker(), new_marker()
key: Final = scenario.key(model_rpm_limit={primary: 1})
assert_served(chat(rig.gateway, key, primary, first), primary, f"chatcmpl-{first}")
assert_refused(chat(rig.gateway, key, primary, second), "model_per_key", primary)
assert rig.traffic() == rig.expected({primary: (first,)})
@pytest.mark.parametrize("raw", ("sometimes", ""), ids=("word", "empty"))
def test_yaml_non_boolean_setting_warns_and_keeps_falling_back(tmp_path: Path, raw: str) -> None:
with fallback_rig(tmp_path, {SETTING: raw}) as rig, rig.gateway.scenario() as scenario:
primary, fallback = rig.groups.primary, rig.groups.fallback
first, second = new_marker(), new_marker()
key: Final = scenario.key(model_rpm_limit={primary: 1})
assert_served(chat(rig.gateway, key, primary, first), primary, f"chatcmpl-{first}")
assert_served(chat(rig.gateway, key, primary, second), fallback, f"chatcmpl-{second}")
assert rig.traffic() == rig.expected({primary: (first,), fallback: (second,)})
eventually(
lambda: rig.log.read_text(errors="replace"),
lambda text: f"general_settings.{SETTING}={raw!r} {NOT_A_BOOLEAN}" in text,
seconds=10,
)
def test_legacy_limiter_per_model_429_skips_fallbacks(tmp_path: Path) -> None:
with fallback_rig(tmp_path, {SETTING: True}, LEGACY_LIMITER) as rig, rig.gateway.scenario() as scenario:
primary: Final = rig.groups.primary
markers: Final = tuple(new_marker() for _ in range(LEGACY_ATTEMPTS))
responses: Final = sent_until_refused(rig, scenario.key(model_rpm_limit={primary: 1}), markers)
*served, refused = responses
assert refused.status_code == 429, tuple(
(response.status_code, response.headers.get("x-litellm-model-group")) for response in responses
)
assert "LiteLLM Rate Limit Handler for rate limit type = model_per_key." in error_message(refused.text)
for marker, response in zip(markers, served):
assert_served(response, primary, f"chatcmpl-{marker}")
assert rig.traffic() == rig.expected({primary: markers[: len(served)]})
def test_database_setting_survives_a_restart_and_refuses_per_model_overflow(tmp_path: Path) -> None:
groups: Final = new_groups()
with ExitStack() as stack:
rig_gateway: Final = stack.enter_context(gateway_from_environment())
database_url: Final = stack.enter_context(scratch_database())
wires: Final = MappingProxyType(
{group: stack.enter_context(wire_server(healthy_reply)) for group in serving_groups(groups)}
)
config: Final = fallback_config(tmp_path, groups, wires, {})
environment: Final = MappingProxyType({"DATABASE_URL": database_url})
replica: Final = ("DATABASE_URL_READ_REPLICA",)
with owned_proxy_process(
rig_gateway, tmp_path, environment, config=config, remove_environment=replica, workers=2
) as writer:
writes: Final = (update_setting(writer.gateway, True), update_setting(writer.gateway, True))
assert tuple(write.status_code for write in writes) == (200, 200), tuple(write.text for write in writes)
restarted: Final = stack.enter_context(
owned_proxy_process(
rig_gateway, tmp_path, environment, config=config, remove_environment=replica, workers=2
)
)
listed: Final = listed_setting(restarted.gateway)
assert [(entry["field_value"], entry["stored_in_db"]) for entry in listed] == [(True, True)], listed
rig: Final = Rig(restarted.gateway, restarted.log, groups, wires)
rig.traffic()
key: Final = string_value(rig.gateway.post("/key/generate", {"model_rpm_limit": {groups.primary: 1}})["key"])
first, second = new_marker(), new_marker()
assert_served(chat(rig.gateway, key, groups.primary, first), groups.primary, f"chatcmpl-{first}")
assert_refused(chat(rig.gateway, key, groups.primary, second), "model_per_key", groups.primary)
assert rig.traffic() == rig.expected({groups.primary: (first,)})

73
tests/integration/sdk/test_interception_converted_stream_options_sdk.py Normal file
View file

@ -0,0 +1,73 @@
from collections.abc import Iterator, Sequence
from typing import Final
import pytest
from integration._support import interception_vendor as iv
from integration._support.wire import Wire, wire_server
from openai.types.responses import WebSearchToolParam
import litellm
from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger
from litellm.responses.streaming_iterator import (
MockResponsesAPIStreamingIterator,
ResponsesAPIStreamingIterator,
SyncResponsesAPIStreamingIterator,
)
from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamingResponse
@pytest.fixture
def wire(monkeypatch: pytest.MonkeyPatch) -> Iterator[Wire]:
monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["azure"])])
with wire_server(iv.respond) as served:
yield served
def _text(events: Sequence[ResponsesAPIStreamingResponse]) -> str:
return "".join(event.delta for event in events if isinstance(event, OutputTextDeltaEvent))
def _assert_one_clean_call(wire: Wire, name: str) -> None:
received: Final = iv.received(wire.drain(), name)
(call,) = received.model_calls
assert call.get("stream") is not True and "stream_options" not in call, call
assert [key for key in call if key.startswith("_")] == [], call
assert received.searches == (), received
def test_sync_responses_stream_with_web_search_keeps_stream_options_off_azure(wire: Wire) -> None:
name: Final = iv.deployment("plain")
stream: Final = litellm.responses(
model=f"azure/{name}",
input=f"Search for {name}",
tools=[WebSearchToolParam(type="web_search")],
stream=True,
stream_options={"include_obfuscation": True},
api_base=wire.url,
api_key=iv.AZURE_KEY,
api_version=iv.API_VERSION,
num_retries=0,
)
assert isinstance(stream, (SyncResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator)), stream
events: Final = tuple(stream)
assert _text(events) == iv.plain(name), events
_assert_one_clean_call(wire, name)
async def test_async_responses_stream_with_web_search_keeps_stream_options_off_azure(wire: Wire) -> None:
name: Final = iv.deployment("plain")
stream: Final = await litellm.aresponses(
model=f"azure/{name}",
input=f"Search for {name}",
tools=[WebSearchToolParam(type="web_search")],
stream=True,
stream_options={"include_obfuscation": True},
api_base=wire.url,
api_key=iv.AZURE_KEY,
api_version=iv.API_VERSION,
num_retries=0,
)
assert isinstance(stream, (ResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator)), stream
events: Final = tuple([event async for event in stream])
assert _text(events) == iv.plain(name), events
_assert_one_clean_call(wire, name)

488
tests/integration/sdk/test_openai_sdk_client_sessions_wire.py Normal file
View file

@ -0,0 +1,488 @@
import asyncio
import base64
import json
import struct
import threading
import zlib
from collections.abc import Awaitable, Callable, Iterator, Mapping
from queue import SimpleQueue
from typing import Final, TypeVar
from urllib.parse import urlsplit
import httpx
import openai
import pytest
from integration._support.vertex import service_account_json
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue
import litellm
from litellm import Router
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.utils import ImageResponse, ModelResponse, TextCompletionResponse
_R: Final = TypeVar("_R")
_PROVIDER_KEY: Final = "sk-scripted-provider"
_API_VERSION: Final = "2024-10-21"
_GATEWAY_PATH: Final = "/gateway.ai.cloudflare.com/v1/scripted-account/scripted-gateway/azure-openai/scripted-resource"
_DEPLOYMENT: Final = "gpt-4o-mini-gateway"
_ANSWER: Final = "wire answer"
_IMAGE_URL: Final = "https://images.example.invalid/variation.png"
def _png_chunk(kind: bytes, data: bytes) -> bytes:
return struct.pack(">I", len(data)) + kind + data + struct.pack(">I", zlib.crc32(kind + data))
_PNG: Final = (
b"\x89PNG\r\n\x1a\n"
+ _png_chunk(b"IHDR", struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0))
+ _png_chunk(b"IDAT", zlib.compress(b"\x00\x00\x00\x00"))
+ _png_chunk(b"IEND", b"")
)
def _json(body: Mapping[str, JsonValue]) -> Reply:
return Reply(body=json.dumps(body).encode())
def _chat_completion() -> Mapping[str, JsonValue]:
return {
"id": "chatcmpl-wire",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": _ANSWER}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
}
def _text_completion() -> Mapping[str, JsonValue]:
return {
"id": "cmpl-wire",
"object": "text_completion",
"created": 1,
"model": "gpt-3.5-turbo-instruct",
"choices": [{"index": 0, "text": _ANSWER, "finish_reason": "stop", "logprobs": None}],
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
}
def _peer(expected_path: str, body: Mapping[str, JsonValue]) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert urlsplit(request.target).path == expected_path, request.target
return _json(body)
return respond
def _held_peer(gate: threading.Event) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert gate.wait(timeout=10), "the timeout cell never released its peer"
return _json(_chat_completion())
return respond
def _only_request(wire: Wire) -> Request:
received: Final = wire.drain()
assert len(received) == 1, received
return received[0]
def _drain(seen: SimpleQueue[str]) -> tuple[str, ...]:
return tuple(seen.get_nowait() for _ in range(seen.qsize()))
def _content(response: ModelResponse) -> str | None:
choice: Final = response.choices[0]
return choice.message.content if isinstance(choice, litellm.Choices) else None
@pytest.fixture
def sync_session(monkeypatch: pytest.MonkeyPatch) -> Iterator[tuple[httpx.Client, SimpleQueue[str]]]:
seen: Final[SimpleQueue[str]] = SimpleQueue()
def record(request: httpx.Request) -> None:
seen.put(str(request.url))
with httpx.Client(event_hooks={"request": [record]}) as client:
monkeypatch.setattr(litellm, "client_session", client)
monkeypatch.setattr(litellm, "aclient_session", None)
yield client, seen
def _run_with_async_session(
monkeypatch: pytest.MonkeyPatch, call: Callable[[], Awaitable[_R]]
) -> tuple[_R, tuple[str, ...]]:
seen: Final[SimpleQueue[str]] = SimpleQueue()
async def record(request: httpx.Request) -> None:
seen.put(str(request.url))
async def run() -> _R:
async with httpx.AsyncClient(event_hooks={"request": [record]}) as client:
monkeypatch.setattr(litellm, "client_session", None)
monkeypatch.setattr(litellm, "aclient_session", client)
return await call()
result: Final = asyncio.run(run())
return result, _drain(seen)
def _deployment(timeout: httpx.Timeout | openai.Timeout, api_base: str, deployment_id: str) -> Router:
return Router(
model_list=[
{
"model_name": "gpt-4o-mini",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_base": api_base,
"api_key": _PROVIDER_KEY,
"timeout": timeout,
"max_retries": 0,
},
"model_info": {"id": deployment_id},
}
],
num_retries=0,
)
def test_f01_async_image_variation_with_only_a_sync_session_set_builds_its_own_async_client(
sync_session: tuple[httpx.Client, SimpleQueue[str]],
) -> None:
with wire_server(_peer("/images/variations", {"created": 1, "data": [{"url": _IMAGE_URL}]})) as wire:
response: Final = asyncio.run(
litellm.aimage_variation(
image=("probe.png", _PNG, "image/png"),
model="dall-e-2",
custom_llm_provider="openai",
api_base=wire.url,
api_key=_PROVIDER_KEY,
num_retries=0,
)
)
assert isinstance(response, ImageResponse), type(response)
assert response.data is not None and response.data[0].url == _IMAGE_URL, response
request: Final = _only_request(wire)
assert request.method == "POST", request.method
assert _PNG in request.body
assert request.headers.get("authorization") == f"Bearer {_PROVIDER_KEY}", request.headers
assert _drain(sync_session[1]) == ()
def test_f02_async_azure_cloudflare_gateway_call_with_only_a_sync_session_set_builds_its_own_async_client(
sync_session: tuple[httpx.Client, SimpleQueue[str]],
) -> None:
with wire_server(_peer(f"{_GATEWAY_PATH}/{_DEPLOYMENT}/chat/completions", _chat_completion())) as wire:
response: Final = asyncio.run(
litellm.acompletion(
model=f"azure/{_DEPLOYMENT}",
messages=[{"role": "user", "content": "via the gateway"}],
api_base=f"{wire.url}{_GATEWAY_PATH}",
api_key=_PROVIDER_KEY,
api_version=_API_VERSION,
num_retries=0,
)
)
assert isinstance(response, ModelResponse), type(response)
assert _content(response) == _ANSWER, response
request: Final = _only_request(wire)
assert urlsplit(request.target).query == f"api-version={_API_VERSION}", request.target
assert request.headers.get("api-key") == _PROVIDER_KEY, request.headers
assert _drain(sync_session[1]) == ()
def test_f03_async_text_completion_uses_the_callers_async_session(monkeypatch: pytest.MonkeyPatch) -> None:
with wire_server(_peer("/completions", _text_completion())) as wire:
response, seen = _run_with_async_session(
monkeypatch,
lambda: litellm.atext_completion(
model="gpt-3.5-turbo-instruct",
prompt="complete this",
custom_llm_provider="openai",
api_base=wire.url,
api_key=_PROVIDER_KEY,
num_retries=0,
),
)
assert isinstance(response, TextCompletionResponse), type(response)
assert response.choices[0].text == _ANSWER, response
assert _only_request(wire).method == "POST"
assert seen == (f"{wire.url}/completions",)
def test_f04_sync_text_completion_uses_the_callers_sync_session(
sync_session: tuple[httpx.Client, SimpleQueue[str]],
) -> None:
with wire_server(_peer("/completions", _text_completion())) as wire:
response: Final = litellm.text_completion(
model="gpt-3.5-turbo-instruct",
prompt="complete this",
custom_llm_provider="openai",
api_base=wire.url,
api_key=_PROVIDER_KEY,
num_retries=0,
)
assert isinstance(response, TextCompletionResponse), type(response)
assert response.choices[0].text == _ANSWER, response
assert _only_request(wire).method == "POST"
assert _drain(sync_session[1]) == (f"{wire.url}/completions",)
def test_f05_sync_moderation_reaches_the_peer_and_returns_its_verdict() -> None:
verdict: Final[Mapping[str, JsonValue]] = {
"id": "modr-wire",
"model": "omni-moderation-latest",
"results": [{"flagged": True, "categories": {"harassment": True}, "category_scores": {"harassment": 0.91}}],
}
with wire_server(_peer("/moderations", verdict)) as wire:
response: Final = litellm.moderation(
input="moderate this",
model="omni-moderation-latest",
api_key=_PROVIDER_KEY,
api_base=wire.url,
)
assert response.results[0].flagged is True, response
request: Final = _only_request(wire)
assert json.loads(request.body) == {"input": "moderate this", "model": "omni-moderation-latest"}, request.body
@pytest.mark.parametrize("timeout_type", (httpx.Timeout, openai.Timeout), ids=("httpx", "openai"))
def test_f06_sync_completion_accepts_a_structured_timeout(
timeout_type: type[httpx.Timeout] | type[openai.Timeout],
) -> None:
with wire_server(_peer("/chat/completions", _chat_completion())) as wire:
response: Final = litellm.completion(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "with a structured timeout"}],
custom_llm_provider="openai",
api_base=wire.url,
api_key=_PROVIDER_KEY,
timeout=timeout_type(5.0, connect=2.0),
num_retries=0,
)
assert isinstance(response, ModelResponse), type(response)
assert _content(response) == _ANSWER, response
assert _only_request(wire).method == "POST"
def test_f07_async_azure_cloudflare_gateway_call_uses_the_callers_async_session(
monkeypatch: pytest.MonkeyPatch,
) -> None:
with wire_server(_peer(f"{_GATEWAY_PATH}/{_DEPLOYMENT}/chat/completions", _chat_completion())) as wire:
response, seen = _run_with_async_session(
monkeypatch,
lambda: litellm.acompletion(
model=f"azure/{_DEPLOYMENT}",
messages=[{"role": "user", "content": "via the gateway"}],
api_base=f"{wire.url}{_GATEWAY_PATH}",
api_key=_PROVIDER_KEY,
api_version=_API_VERSION,
num_retries=0,
),
)
assert isinstance(response, ModelResponse), type(response)
assert _content(response) == _ANSWER, response
assert _only_request(wire).headers.get("api-key") == _PROVIDER_KEY
assert seen == (f"{wire.url}{_GATEWAY_PATH}/{_DEPLOYMENT}/chat/completions?api-version={_API_VERSION}",)
def test_f08_sync_image_variation_uses_the_callers_sync_session(
sync_session: tuple[httpx.Client, SimpleQueue[str]],
) -> None:
with wire_server(_peer("/images/variations", {"created": 1, "data": [{"url": _IMAGE_URL}]})) as wire:
response: Final = litellm.image_variation(
image=("probe.png", _PNG, "image/png"),
model="dall-e-2",
custom_llm_provider="openai",
api_base=wire.url,
api_key=_PROVIDER_KEY,
num_retries=0,
)
assert isinstance(response, ImageResponse), type(response)
assert response.data is not None and response.data[0].url == _IMAGE_URL, response
assert _PNG in _only_request(wire).body
assert _drain(sync_session[1]) == (f"{wire.url}/images/variations",)
@pytest.mark.parametrize("timeout_type", (httpx.Timeout, openai.Timeout), ids=("httpx", "openai"))
def test_f09_router_deployment_keeps_a_structured_timeout_as_httpx_and_serves(
timeout_type: type[httpx.Timeout] | type[openai.Timeout],
) -> None:
with wire_server(_peer("/chat/completions", _chat_completion())) as wire:
router: Final = _deployment(timeout_type(5.0, connect=2.0), wire.url, "deployment-f09")
response: Final = router.completion(
model="gpt-4o-mini", messages=[{"role": "user", "content": "through the router"}]
)
assert isinstance(response, ModelResponse), type(response)
assert _content(response) == _ANSWER, response
assert _only_request(wire).method == "POST"
deployment: Final = router.get_deployment("deployment-f09")
assert deployment is not None
stored: Final = deployment.litellm_params.timeout
assert isinstance(stored, httpx.Timeout), type(stored)
assert (stored.read, stored.connect) == (5.0, 2.0), stored
def test_f10_router_deployment_structured_read_timeout_fires_once_as_a_408() -> None:
gate: Final = threading.Event()
with wire_server(_held_peer(gate)) as wire:
router: Final = _deployment(httpx.Timeout(0.5, connect=2.0), wire.url, "deployment-f10")
with pytest.raises(litellm.Timeout) as raised:
router.completion(model="gpt-4o-mini", messages=[{"role": "user", "content": "held upstream"}])
gate.set()
assert raised.value.status_code == 408, raised.value
assert len(wire.drain()) == 1
_SPEECH_AUDIO: Final = b"OggS" + bytes(range(60))
_SPEECH_AUDIO_REPLY: Final = Reply(body=_SPEECH_AUDIO, content_type="audio/ogg")
_ELEVENLABS_VOICE: Final = "21m00Tcm4TlvDq8ikWAM"
_VERTEX_PROJECT: Final = "scripted-project"
_SPEECH_ROUTES: Final[Mapping[str, tuple[Mapping[str, str], str, Reply]]] = {
"openai": (
{"model": "openai/gpt-4o-mini-tts", "voice": "alloy", "api_key": _PROVIDER_KEY},
"/audio/speech",
_SPEECH_AUDIO_REPLY,
),
"azure": (
{"model": "azure/tts-deployment", "voice": "alloy", "api_key": _PROVIDER_KEY, "api_version": _API_VERSION},
"/openai/deployments/tts-deployment/audio/speech",
_SPEECH_AUDIO_REPLY,
),
"azure_ava": (
{"model": "azure/speech/azure-tts", "voice": "alloy", "api_key": _PROVIDER_KEY},
"/cognitiveservices/v1",
_SPEECH_AUDIO_REPLY,
),
"elevenlabs": (
{"model": "elevenlabs/eleven_multilingual_v2", "voice": _ELEVENLABS_VOICE, "api_key": _PROVIDER_KEY},
f"/v1/text-to-speech/{_ELEVENLABS_VOICE}",
_SPEECH_AUDIO_REPLY,
),
"edenai": (
{"model": "edenai/openai", "voice": "alloy", "api_key": _PROVIDER_KEY},
"/audio/speech",
_SPEECH_AUDIO_REPLY,
),
"minimax": (
{"model": "minimax/speech-02-hd", "voice": "alloy", "api_key": _PROVIDER_KEY},
"/v1/t2a_v2",
_json({"data": {"audio": _SPEECH_AUDIO.hex()}, "base_resp": {"status_code": 0, "status_msg": "success"}}),
),
"mistral": (
{"model": "mistral/voxtral-mini-tts-2603", "voice": "alloy", "api_key": _PROVIDER_KEY},
"/v1/audio/speech",
_json({"audio_data": base64.b64encode(_SPEECH_AUDIO).decode()}),
),
"aws_polly": (
{
"model": "aws_polly/neural",
"voice": "Joanna",
"aws_access_key_id": "AKIASCRIPTED",
"aws_secret_access_key": "scripted-secret",
"aws_region_name": "us-east-1",
},
"/v1/speech",
_SPEECH_AUDIO_REPLY,
),
}
_PAYMENT_REQUIRED_MODELS: Final[Mapping[str, type[litellm.BadRequestError]]] = {
"openai/gpt-4o-mini": litellm.BadRequestError,
"anthropic/claude-sonnet-4-5": litellm.PaymentRequiredError,
}
def _speech_peer(expected_path: str, reply: Reply) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert urlsplit(request.target).path == expected_path, request.target
return reply
return respond
def _vertex_speech_peer(request: Request) -> Reply:
if urlsplit(request.target).path == "/_oauth/token":
return _json({"access_token": "scripted-vertex-token", "expires_in": 3600, "token_type": "Bearer"})
return _json({"audioContent": base64.b64encode(_SPEECH_AUDIO).decode()})
def _payment_required_peer(request: Request) -> Reply:
return Reply(
status=402,
body=json.dumps({"error": {"message": "scripted 402", "type": "invalid_request_error"}}).encode(),
)
@pytest.mark.parametrize("provider", tuple(_SPEECH_ROUTES))
def test_f11_sync_speech_accepts_an_sdk_timeout_on_every_provider_branch(provider: str) -> None:
route, expected_path, reply = _SPEECH_ROUTES[provider]
with wire_server(_speech_peer(expected_path, reply)) as wire:
response: Final = litellm.speech(
input="speak this",
api_base=wire.url,
timeout=openai.Timeout(5.0, connect=2.0),
**route,
)
assert response.content == _SPEECH_AUDIO, response.content[:16]
assert _only_request(wire).method == "POST"
def test_f13_sync_vertex_speech_accepts_an_sdk_timeout() -> None:
with wire_server(_vertex_speech_peer) as wire:
response: Final = litellm.speech(
model="vertex_ai/chirp",
voice="alloy",
input="speak this",
api_base=wire.url,
vertex_credentials=service_account_json(_VERTEX_PROJECT, wire.url),
vertex_project=_VERTEX_PROJECT,
vertex_location="us-central1",
timeout=openai.Timeout(5.0, connect=2.0),
)
assert response.content == _SPEECH_AUDIO, response.content[:16]
assert tuple((seen.method, urlsplit(seen.target).path) for seen in wire.drain()) == (
("POST", "/_oauth/token"),
("POST", "/"),
)
def _runwayml_rejecting_peer(request: Request) -> Reply:
assert urlsplit(request.target).path == "/v1/text_to_speech", request.target
return Reply(status=400, body=json.dumps({"error": "scripted 400"}).encode())
def test_f14_sync_runwayml_speech_sends_its_task_with_an_sdk_timeout() -> None:
with wire_server(_runwayml_rejecting_peer) as wire:
with pytest.raises(BaseLLMException) as raised:
litellm.speech(
model="runwayml/eleven_multilingual_v2",
voice="Maya",
input="speak this",
api_base=wire.url,
api_key=_PROVIDER_KEY,
timeout=openai.Timeout(5.0, connect=2.0),
)
assert raised.value.status_code == 400, raised.value
assert "scripted 400" in raised.value.message, raised.value.message
assert _only_request(wire).method == "POST"
@pytest.mark.parametrize("model", tuple(_PAYMENT_REQUIRED_MODELS))
def test_f12_sync_completion_maps_an_upstream_402_with_its_message(model: str) -> None:
with wire_server(_payment_required_peer) as wire:
with pytest.raises(_PAYMENT_REQUIRED_MODELS[model]) as raised:
litellm.completion(
model=model,
messages=[{"role": "user", "content": "out of credit"}],
api_base=wire.url,
api_key=_PROVIDER_KEY,
num_retries=0,
max_retries=0,
)
assert raised.value.status_code == 402, raised.value
assert "scripted 402" in raised.value.message, raised.value.message
assert len(wire.drain()) == 1

905
tests/integration/streaming/test_interception_converted_stream_options.py Normal file
View file

@ -0,0 +1,905 @@
import asyncio
import json
import signal
import threading
import uuid
from collections.abc import Callable, Generator, Iterator, Mapping, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from queue import SimpleQueue
from types import MappingProxyType
from typing import Final
from urllib.parse import urlsplit
import anthropic
import httpx
import openai
import psutil
import pytest
import yaml
from anthropic.types import RawContentBlockDeltaEvent, TextDelta, WebSearchTool20250305Param
from integration._support import interception_vendor as iv
from integration._support import responses_vendor as rv
from integration._support.client import (
Gateway,
Scenario,
eventually,
gateway_from_environment,
object_value,
string_value,
)
from integration._support.database import read_rows
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from openai.types.chat import ChatCompletionChunk, ChatCompletionFunctionToolParam, ChatCompletionUserMessageParam
from openai.types.responses import ResponseCompletedEvent, ResponseTextDeltaEvent, WebSearchToolParam
from pydantic import JsonValue
_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120
pytestmark: Final = pytest.mark.timeout(_CELL_SECONDS)
_PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
_NO_CACHE: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}})
_USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType({"include_usage": True})
_OBFUSCATION: Final[Mapping[str, JsonValue]] = MappingProxyType({"include_obfuscation": True})
_PARAMETERS: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}
)
_SEARCH_FUNCTION: Final[Mapping[str, JsonValue]] = MappingProxyType(
{
"type": "function",
"function": {"name": iv.SEARCH_TOOL, "description": "Search the web", "parameters": dict(_PARAMETERS)},
}
)
_CODE_INTERPRETER: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"type": "code_interpreter", "container": {"type": "auto"}}
)
_HEADROOM_RETRIEVE: Final[Mapping[str, JsonValue]] = MappingProxyType(
{
"type": "function",
"function": {"name": "headroom_retrieve", "description": "Retrieve content", "parameters": dict(_PARAMETERS)},
}
)
_RESPONSES_SEARCH: Final[Mapping[str, JsonValue]] = MappingProxyType({"type": "web_search"})
_BURST_MODEL: Final = f"interception-burst-{uuid.uuid4().hex}"
_BURST_DEPLOYMENT: Final = iv.deployment("search")
_STARTED_WORKER: Final = "Started server process ["
_STARTUP_COMPLETE: Final = "Application startup complete."
_MALFORMED: Final = (
pytest.param(5, id="int"),
pytest.param([{"include_usage": True}], id="list"),
pytest.param("", id="empty-string"),
pytest.param("x" * 5120, id="5kb-string"),
)
@dataclass(frozen=True, slots=True)
class _Rig:
wire: Wire
proxy: OwnedProxy
@property
def gateway(self) -> Gateway:
return self.proxy.gateway
@property
def base_url(self) -> str:
return str(self.proxy.gateway.client.base_url).rstrip("/")
@dataclass(frozen=True, slots=True)
class _Deployment:
model: str
name: str
@dataclass(frozen=True, slots=True)
class _Stream:
status: int
text: str
@property
def frames(self) -> tuple[Mapping[str, JsonValue], ...]:
return tuple(
iv.JSON_OBJECT.validate_json(line[6:]) for line in self.text.splitlines() if line.startswith("data: {")
)
@property
def content(self) -> str:
return "".join(_delta_text(frame) for frame in self.frames)
@property
def ids(self) -> frozenset[str]:
return frozenset(string_value(frame["id"]) for frame in self.frames)
@property
def usages(self) -> tuple[Mapping[str, JsonValue], ...]:
return tuple(object_value(frame["usage"]) for frame in self.frames if frame.get("usage") is not None)
@property
def deltas(self) -> str:
return "".join(
string_value(frame["delta"]) for frame in self.frames if frame.get("type") == "response.output_text.delta"
)
@property
def completed(self) -> Mapping[str, JsonValue]:
(frame,) = [frame for frame in self.frames if frame.get("type") == "response.completed"]
return object_value(frame["response"])
def _delta_text(frame: Mapping[str, JsonValue]) -> str:
choices: Final = rv.ITEMS.validate_python(frame.get("choices") or [])
return "".join(_text_of(object_value(choice.get("delta") or {})) for choice in choices)
def _text_of(delta: Mapping[str, JsonValue]) -> str:
content: Final = delta.get("content")
return content if isinstance(content, str) else ""
def _config(
directory: Path,
wire: Wire,
general_settings: Mapping[str, JsonValue],
models: Sequence[Mapping[str, JsonValue]] = (),
) -> Path:
base: Final = object_value(yaml.safe_load(_PROXY_CONFIG.read_text()))
config: Final = {
**base,
"model_list": [*rv.ITEMS.validate_python(base["model_list"]), *(dict(model) for model in models)],
"general_settings": {**object_value(base["general_settings"]), **general_settings},
"litellm_settings": {
**object_value(base["litellm_settings"]),
"callbacks": ["websearch_interception", "code_interpreter_interception"],
"websearch_interception_params": {"enabled_providers": ["azure"], "search_tool_name": "integration-search"},
"code_interpreter_interception_params": {"enabled_providers": ["azure"]},
},
"search_tools": [
{
"search_tool_name": "integration-search",
"litellm_params": {
"search_provider": "tavily",
"api_key": iv.TAVILY_KEY,
"api_base": f"{wire.url}/tavily",
},
}
],
"guardrails": [
{
"guardrail_name": "integration-headroom",
"litellm_params": {
"guardrail": "headroom",
"mode": "post_call",
"api_base": "http://127.0.0.1:9",
"unreachable_fallback": "fail_open",
"default_on": False,
},
}
],
}
path: Final = directory / f"interception-{uuid.uuid4().hex}.yaml"
path.write_text(yaml.safe_dump(config))
return path
@contextmanager
def _interception_proxy(
directory: Path,
respond: Callable[[Request], Reply],
general_settings: Mapping[str, JsonValue],
models: Callable[[Wire], Sequence[Mapping[str, JsonValue]]] = lambda _: (),
) -> Generator[_Rig, None, None]:
with gateway_from_environment() as gateway, wire_server(respond) as wire:
config: Final = _config(directory, wire, general_settings, models(wire))
with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned:
yield _Rig(wire, owned)
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
with _interception_proxy(tmp_path_factory.mktemp("interception-rig"), iv.respond, {}) as built:
yield built
@pytest.fixture
def scenario(rig: _Rig) -> Iterator[Scenario]:
rig.wire.drain()
with rig.gateway.scenario() as opened:
yield opened
def _deploy(rig: _Rig, scenario: Scenario, mode: iv.Mode, *, api_key: str = iv.AZURE_KEY) -> _Deployment:
name: Final = iv.deployment(mode)
model: Final = scenario.model(
model=f"azure/{name}", api_base=rig.wire.url, api_key=api_key, api_version=iv.API_VERSION
)
return _Deployment(model, name)
def _chat_body(
deployment: _Deployment,
*,
stream: bool = True,
tool: Mapping[str, JsonValue] = _SEARCH_FUNCTION,
extra: Mapping[str, JsonValue] = _NO_CACHE,
) -> Mapping[str, JsonValue]:
return {
"model": deployment.model,
"messages": [{"role": "user", "content": f"Search for {deployment.name}"}],
"tools": [dict(tool)],
"stream": stream,
**extra,
}
def _responses_body(deployment: _Deployment, extra: Mapping[str, JsonValue] = _NO_CACHE) -> Mapping[str, JsonValue]:
return {
"model": deployment.model,
"input": f"Search for {deployment.name}",
"tools": [dict(_RESPONSES_SEARCH)],
"stream": True,
**extra,
}
def _post(rig: _Rig, path: str, body: Mapping[str, JsonValue], *, key: str | None = None) -> _Stream:
response: Final = rig.gateway.request("POST", path, body, key=key)
return _Stream(response.status_code, response.text)
def _chat(rig: _Rig, body: Mapping[str, JsonValue]) -> _Stream:
return _post(rig, "/v1/chat/completions", body)
def _responses(rig: _Rig, body: Mapping[str, JsonValue]) -> _Stream:
return _post(rig, "/v1/responses", body)
def _upstream(rig: _Rig, deployment: _Deployment) -> iv.Received:
return iv.received(rig.wire.drain(), deployment.name)
def _leaks(call: Mapping[str, JsonValue]) -> bool:
return "stream_options" in call or call.get("stream") is True or any(key.startswith("_") for key in call)
def _assert_upstream(received: iv.Received, *, model_calls: int, searches: int) -> None:
assert len(received.model_calls) == model_calls, received
assert len(received.searches) == searches, received
assert [call for call in received.model_calls if _leaks(call)] == [], received.model_calls
def _totals(usage: Mapping[str, JsonValue]) -> int:
total: Final = usage["total_tokens"]
assert isinstance(total, int), usage
return total
def _assert_restreamed(stream: _Stream, text: str, *, usage: Mapping[str, int] | None) -> str:
assert stream.status == 200, stream.text
assert stream.content == text, stream.text
assert [_totals(item) for item in stream.usages] == ([] if usage is None else [usage["total_tokens"]]), stream.text
(identity,) = stream.ids
return identity
def _assert_responses_restreamed(stream: _Stream, text: str) -> str:
assert stream.status == 200, stream.text
assert stream.deltas == text, stream.text
completed: Final = stream.completed
assert _totals(object_value(completed["usage"])) == iv.RESPONSES_FOLLOWUP_USAGE["total_tokens"], completed
return string_value(completed["id"])
def _spend_rows(model: str, expected: int) -> tuple[Mapping[str, JsonValue], ...]:
return tuple(
eventually(
lambda: read_rows(
'SELECT request_id, status, cache_hit FROM "LiteLLM_SpendLogs" WHERE model_group = %s', (model,)
),
lambda rows: sum(1 for row in rows if row["status"] == "success") >= expected,
seconds=70,
)
)
def _assert_landed_once(model: str, identities: frozenset[str], *, rows: int) -> tuple[Mapping[str, JsonValue], ...]:
landed: Final = _spend_rows(model, rows)
assert len(landed) == rows, landed
assert frozenset(string_value(row["request_id"]) for row in landed if row["status"] == "success") == identities, (
landed
)
return landed
def _assert_response_landed_once(model: str, identity: str) -> None:
(row,) = _spend_rows(model, 1)
assert row["status"] == "success" and rv.same_response(string_value(row["request_id"]), identity), row
def _openai(rig: _Rig) -> openai.OpenAI:
return openai.OpenAI(base_url=f"{rig.base_url}/v1", api_key=rig.gateway.key, max_retries=0)
def _async_openai(rig: _Rig) -> openai.AsyncOpenAI:
return openai.AsyncOpenAI(base_url=f"{rig.base_url}/v1", api_key=rig.gateway.key, max_retries=0)
def _sdk_tool() -> ChatCompletionFunctionToolParam:
return ChatCompletionFunctionToolParam(
type="function",
function={"name": iv.SEARCH_TOOL, "description": "Search the web", "parameters": dict(_PARAMETERS)},
)
def _sdk_message(deployment: _Deployment) -> ChatCompletionUserMessageParam:
return ChatCompletionUserMessageParam(role="user", content=f"Search for {deployment.name}")
def _chunk_text(chunk: ChatCompletionChunk) -> str:
return "".join(choice.delta.content or "" for choice in chunk.choices)
def _assert_sdk_chunks(chunks: Sequence[ChatCompletionChunk], deployment: _Deployment) -> str:
assert "".join(map(_chunk_text, chunks)) == iv.answer(deployment.name), chunks
usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None)
assert [usage.total_tokens for usage in usages] == [iv.FOLLOWUP_USAGE["total_tokens"]], chunks
(identity,) = {chunk.id for chunk in chunks}
return identity
def _assert_search_turn(rig: _Rig, deployment: _Deployment, identity: str) -> None:
received: Final = _upstream(rig, deployment)
_assert_upstream(received, model_calls=2, searches=1)
assert [search["query"] for search in received.searches] == [iv.query(deployment.name)], received.searches
_assert_landed_once(deployment.model, frozenset({identity}), rows=1)
def test_openai_sdk_chat_stream_through_web_search_is_restreamed_with_its_usage_chunk(
rig: _Rig, scenario: Scenario
) -> None:
deployment: Final = _deploy(rig, scenario, "search")
chunks: Final = tuple(
_openai(rig).chat.completions.create(
model=deployment.model,
messages=[_sdk_message(deployment)],
tools=[_sdk_tool()],
stream=True,
stream_options={"include_usage": True},
extra_body=dict(_NO_CACHE),
)
)
_assert_search_turn(rig, deployment, _assert_sdk_chunks(chunks, deployment))
async def test_async_openai_sdk_chat_stream_through_web_search_is_restreamed_with_its_usage_chunk(
rig: _Rig, scenario: Scenario
) -> None:
deployment: Final = _deploy(rig, scenario, "search")
stream: Final = await _async_openai(rig).chat.completions.create(
model=deployment.model,
messages=[_sdk_message(deployment)],
tools=[_sdk_tool()],
stream=True,
stream_options={"include_usage": True},
extra_body=dict(_NO_CACHE),
)
chunks: Final = tuple([chunk async for chunk in stream])
_assert_search_turn(rig, deployment, _assert_sdk_chunks(chunks, deployment))
def test_raw_chat_stream_without_stream_options_answers_without_a_usage_chunk(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "search")
identity: Final = _assert_restreamed(_chat(rig, _chat_body(deployment)), iv.answer(deployment.name), usage=None)
_assert_search_turn(rig, deployment, identity)
def test_openai_sdk_responses_stream_through_web_search_keeps_stream_options_off_azure(
rig: _Rig, scenario: Scenario
) -> None:
deployment: Final = _deploy(rig, scenario, "search")
events: Final = tuple(
_openai(rig).responses.create(
model=deployment.model,
input=f"Search for {deployment.name}",
tools=[WebSearchToolParam(type="web_search")],
stream=True,
stream_options={"include_obfuscation": True},
extra_body=dict(_NO_CACHE),
)
)
text: Final = "".join(event.delta for event in events if isinstance(event, ResponseTextDeltaEvent))
(completed,) = [event for event in events if isinstance(event, ResponseCompletedEvent)]
assert text == iv.answer(deployment.name), events
usage: Final = completed.response.usage
assert usage is not None and usage.total_tokens == iv.RESPONSES_FOLLOWUP_USAGE["total_tokens"], completed
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
_assert_response_landed_once(deployment.model, completed.response.id)
async def test_async_openai_sdk_responses_stream_through_web_search_keeps_stream_options_off_azure(
rig: _Rig, scenario: Scenario
) -> None:
deployment: Final = _deploy(rig, scenario, "search")
stream: Final = await _async_openai(rig).responses.create(
model=deployment.model,
input=f"Search for {deployment.name}",
tools=[WebSearchToolParam(type="web_search")],
stream=True,
stream_options={"include_obfuscation": True},
extra_body=dict(_NO_CACHE),
)
events: Final = tuple([event async for event in stream])
text: Final = "".join(event.delta for event in events if isinstance(event, ResponseTextDeltaEvent))
(completed,) = [event for event in events if isinstance(event, ResponseCompletedEvent)]
assert text == iv.answer(deployment.name), events
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
_assert_response_landed_once(deployment.model, completed.response.id)
def test_raw_responses_stream_without_stream_options_answers(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "search")
identity: Final = _assert_responses_restreamed(
_responses(rig, _responses_body(deployment)), iv.answer(deployment.name)
)
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
_assert_response_landed_once(deployment.model, identity)
def test_openai_sdk_chat_without_streaming_through_web_search_answers_in_json(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "search")
completion: Final = _openai(rig).chat.completions.create(
model=deployment.model,
messages=[_sdk_message(deployment)],
tools=[_sdk_tool()],
extra_body=dict(_NO_CACHE),
)
assert completion.choices[0].message.content == iv.answer(deployment.name), completion
assert completion.usage is not None, completion
_assert_search_turn(rig, deployment, completion.id)
@pytest.mark.parametrize(
("tool", "mode"),
(
pytest.param(_CODE_INTERPRETER, "plain", id="code-interpreter"),
pytest.param(_HEADROOM_RETRIEVE, "plain", id="headroom-retrieve"),
pytest.param(_SEARCH_FUNCTION, "plain", id="web-search-unused"),
),
)
def test_converted_chat_stream_with_usage_replays_the_single_call_usage(
rig: _Rig, scenario: Scenario, tool: Mapping[str, JsonValue], mode: iv.Mode
) -> None:
deployment: Final = _deploy(rig, scenario, mode)
stream: Final = _chat(rig, _chat_body(deployment, tool=tool, extra={**_NO_CACHE, "stream_options": dict(_USAGE)}))
identity: Final = _assert_restreamed(stream, iv.plain(deployment.name), usage=iv.FIRST_USAGE)
_assert_upstream(_upstream(rig, deployment), model_calls=1, searches=0)
_assert_landed_once(deployment.model, frozenset({identity}), rows=1)
def test_anthropic_sdk_messages_stream_with_only_web_search_on_azure_answers_from_the_search(
rig: _Rig, scenario: Scenario
) -> None:
deployment: Final = _deploy(rig, scenario, "search")
prompt: Final = f"Search for {deployment.name}"
client: Final = anthropic.Anthropic(base_url=rig.base_url, api_key=rig.gateway.key, max_retries=0)
events: Final = tuple(
client.messages.create(
model=deployment.model,
max_tokens=64,
messages=[{"role": "user", "content": prompt}],
tools=[WebSearchTool20250305Param(type="web_search_20250305", name="web_search")],
stream=True,
extra_body=dict(_NO_CACHE),
)
)
deltas: Final = tuple(event.delta for event in events if isinstance(event, RawContentBlockDeltaEvent))
text: Final = "".join(delta.text for delta in deltas if isinstance(delta, TextDelta))
assert text == f"Title: title-{prompt}\nURL: https://example.test/result\nSnippet: snippet-{prompt}", events
received: Final = _upstream(rig, deployment)
_assert_upstream(received, model_calls=0, searches=1)
assert [search.get("query") for search in received.searches] == [prompt], received
def test_per_request_fallback_from_a_missing_deployment_restreams_with_usage(rig: _Rig, scenario: Scenario) -> None:
missing: Final = _deploy(rig, scenario, "missing")
good: Final = _deploy(rig, scenario, "search")
body: Final = _chat_body(missing, extra={**_NO_CACHE, "stream_options": dict(_USAGE), "fallbacks": [good.model]})
stream: Final = _chat(rig, {**body, "messages": [{"role": "user", "content": "Search the web"}]})
identity: Final = _assert_restreamed(stream, iv.answer(good.name), usage=iv.FOLLOWUP_USAGE)
drained: Final = rig.wire.drain()
_assert_upstream(iv.received(drained, missing.name), model_calls=1, searches=0)
_assert_upstream(iv.received(drained, good.name), model_calls=2, searches=1)
_assert_landed_once(good.model, frozenset({identity}), rows=1)
def test_chat_stream_cache_hit_replays_the_answer_without_calling_azure(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "search")
body: Final = _chat_body(deployment, extra={"stream_options": dict(_USAGE)})
first: Final = _assert_restreamed(_chat(rig, body), iv.answer(deployment.name), usage=iv.FOLLOWUP_USAGE)
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
second: Final = _chat(rig, body)
assert second.status == 200 and second.content == iv.answer(deployment.name), second.text
_assert_upstream(_upstream(rig, deployment), model_calls=0, searches=0)
rows: Final = _spend_rows(deployment.model, 2)
served: Final = tuple(string_value(row["request_id"]) for row in rows if str(row["cache_hit"]) != "True")
cached: Final = tuple(string_value(row["request_id"]) for row in rows if str(row["cache_hit"]) == "True")
assert served == (first,) and len(cached) == 1 and cached[0].startswith(first), rows
def test_responses_stream_cache_hit_replays_the_answer_without_calling_azure(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "search")
body: Final = _responses_body(deployment, {"stream_options": dict(_OBFUSCATION)})
_assert_responses_restreamed(_responses(rig, body), iv.answer(deployment.name))
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
second: Final = _responses(rig, body)
assert second.status == 200 and second.deltas == iv.answer(deployment.name), second.text
_assert_upstream(_upstream(rig, deployment), model_calls=0, searches=0)
def _with_duplicated_stream_options(body: Mapping[str, JsonValue]) -> bytes:
return (json.dumps({**body, "stream_options": dict(_USAGE)})[:-1] + ', "stream_options": 5}').encode()
@pytest.mark.parametrize("stream_options", _MALFORMED)
def test_malformed_chat_stream_options_never_reach_azure(
rig: _Rig, scenario: Scenario, stream_options: JsonValue
) -> None:
deployment: Final = _deploy(rig, scenario, "search")
stream: Final = _chat(rig, _chat_body(deployment, extra={**_NO_CACHE, "stream_options": stream_options}))
identity: Final = _assert_restreamed(stream, iv.answer(deployment.name), usage=None)
_assert_search_turn(rig, deployment, identity)
def test_duplicated_chat_stream_options_key_never_reaches_azure(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "search")
response: Final = rig.gateway.client.post(
"/v1/chat/completions",
content=_with_duplicated_stream_options(_chat_body(deployment)),
headers={"Authorization": f"Bearer {rig.gateway.key}", "Content-Type": "application/json"},
)
identity: Final = _assert_restreamed(
_Stream(response.status_code, response.text), iv.answer(deployment.name), usage=None
)
_assert_search_turn(rig, deployment, identity)
@pytest.mark.parametrize("stream_options", _MALFORMED)
def test_malformed_responses_stream_options_never_reach_azure(
rig: _Rig, scenario: Scenario, stream_options: JsonValue
) -> None:
deployment: Final = _deploy(rig, scenario, "search")
stream: Final = _responses(rig, _responses_body(deployment, {**_NO_CACHE, "stream_options": stream_options}))
identity: Final = _assert_responses_restreamed(stream, iv.answer(deployment.name))
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
_assert_response_landed_once(deployment.model, identity)
def test_forged_interception_keys_on_a_json_request_never_reach_azure(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "search")
forged: Final = {
**_NO_CACHE,
"_websearch_interception_stream_options": dict(_USAGE),
"_websearch_interception_converted_stream": True,
"_headroom_interception_stream_options": dict(_USAGE),
"_code_interpreter_interception_stream_options": dict(_USAGE),
"max_agentic_loops": 50,
}
response: Final = rig.gateway.request(
"POST", "/v1/chat/completions", _chat_body(deployment, stream=False, extra=forged)
)
assert response.status_code == 200, response.text
assert response.headers["content-type"].startswith("application/json"), response.headers
body: Final = iv.JSON_OBJECT.validate_json(response.content)
(choice,) = rv.ITEMS.validate_python(body["choices"])
assert object_value(choice["message"])["content"] == iv.answer(deployment.name), body
_assert_search_turn(rig, deployment, string_value(body["id"]))
def test_forged_garbage_stash_on_a_stream_is_replaced_by_the_real_stream_options(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "plain")
extra: Final = {**_NO_CACHE, "stream_options": dict(_USAGE), "_code_interpreter_interception_stream_options": "x"}
stream: Final = _chat(rig, _chat_body(deployment, tool=_CODE_INTERPRETER, extra=extra))
identity: Final = _assert_restreamed(stream, iv.plain(deployment.name), usage=iv.FIRST_USAGE)
_assert_upstream(_upstream(rig, deployment), model_calls=1, searches=0)
_assert_landed_once(deployment.model, frozenset({identity}), rows=1)
def test_wrong_azure_key_surfaces_azures_401(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "search", api_key="wrong-azure-key")
stream: Final = _chat(rig, _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)}))
assert stream.status == 401, stream.text
assert "Access denied due to invalid subscription key" in stream.text, stream.text
assert len(_upstream(rig, deployment).model_calls) == 1
def _tool_call(call: Mapping[str, JsonValue]) -> tuple[str, Mapping[str, JsonValue]]:
function: Final = object_value(call["function"])
return string_value(function["name"]), iv.JSON_OBJECT.validate_json(string_value(function["arguments"]))
def _stream_choices(stream: _Stream) -> Iterator[Mapping[str, JsonValue]]:
for frame in stream.frames:
yield from rv.ITEMS.validate_python(frame.get("choices") or [])
def _streamed_tool_calls(stream: _Stream) -> Iterator[tuple[str, Mapping[str, JsonValue]]]:
for choice in _stream_choices(stream):
yield from map(
_tool_call, rv.ITEMS.validate_python(object_value(choice.get("delta") or {}).get("tool_calls") or [])
)
def _handed_back(deployment: _Deployment) -> list[tuple[str, Mapping[str, JsonValue]]]:
return [(iv.SEARCH_TOOL, {"query": iv.query(deployment.name)})]
def _assert_succeeded_once(model: str, identity: str) -> None:
rows: Final = _spend_rows(model, 1)
assert [string_value(row["request_id"]) for row in rows if row["status"] == "success"] == [identity], rows
def test_failed_follow_up_call_hands_the_search_tool_call_back_with_its_usage(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "followup-fails")
stream: Final = _chat(rig, _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)}))
assert stream.status == 200 and stream.content == "", stream.text
assert list(_streamed_tool_calls(stream)) == _handed_back(deployment), stream.text
finished: Final = [choice["finish_reason"] for choice in _stream_choices(stream) if choice.get("finish_reason")]
assert finished == ["tool_calls"], stream.text
assert [_totals(usage) for usage in stream.usages] == [iv.FIRST_USAGE["total_tokens"]], stream.text
(identity,) = stream.ids
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
_assert_succeeded_once(deployment.model, identity)
assert rig.gateway.chat(_deploy(rig, scenario, "plain").model)["object"] == "chat.completion"
def test_failed_follow_up_call_on_a_json_request_hands_the_search_tool_call_back(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "followup-fails")
response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(deployment, stream=False))
assert response.status_code == 200, response.text
body: Final = iv.JSON_OBJECT.validate_json(response.content)
(choice,) = rv.ITEMS.validate_python(body["choices"])
calls: Final = rv.ITEMS.validate_python(object_value(choice["message"])["tool_calls"])
assert [_tool_call(call) for call in calls] == _handed_back(deployment), body
assert choice["finish_reason"] == "tool_calls", body
assert _totals(object_value(body["usage"])) == iv.FIRST_USAGE["total_tokens"], body
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
_assert_succeeded_once(deployment.model, string_value(body["id"]))
def test_search_outage_still_restreams_the_follow_up_answer(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "search-down")
stream: Final = _chat(rig, _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)}))
identity: Final = _assert_restreamed(stream, iv.answer(deployment.name), usage=iv.FOLLOWUP_USAGE)
_assert_search_turn(rig, deployment, identity)
def test_unknown_virtual_key_is_refused_before_azure(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "search")
body: Final = _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)})
stream: Final = _post(rig, "/v1/chat/completions", body, key=f"sk-{uuid.uuid4().hex}")
assert stream.status == 401, stream.text
assert _upstream(rig, deployment).model_calls == ()
def test_json_request_with_its_own_stream_options_still_gets_azures_refusal(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "search")
body: Final = _chat_body(deployment, stream=False, extra={**_NO_CACHE, "stream_options": dict(_USAGE)})
response: Final = rig.gateway.request("POST", "/v1/chat/completions", body)
assert response.status_code == 400, response.text
assert iv.STREAM_OPTIONS_REFUSAL in response.text, response.text
assert [call.get("stream_options") for call in _upstream(rig, deployment).model_calls] == [dict(_USAGE)]
@pytest.mark.parametrize(
"stream_options",
(
pytest.param({"include_usage": False}, id="include-usage-false"),
pytest.param(None, id="null"),
pytest.param({}, id="empty"),
),
)
def test_chat_stream_without_a_usage_request_answers_without_a_usage_chunk(
rig: _Rig, scenario: Scenario, stream_options: JsonValue
) -> None:
deployment: Final = _deploy(rig, scenario, "search")
stream: Final = _chat(rig, _chat_body(deployment, extra={**_NO_CACHE, "stream_options": stream_options}))
identity: Final = _assert_restreamed(stream, iv.answer(deployment.name), usage=None)
_assert_search_turn(rig, deployment, identity)
@pytest.mark.parametrize("stream_options", (pytest.param(None, id="null"), pytest.param({}, id="empty")))
def test_responses_stream_with_empty_stream_options_answers(
rig: _Rig, scenario: Scenario, stream_options: JsonValue
) -> None:
deployment: Final = _deploy(rig, scenario, "search")
stream: Final = _responses(rig, _responses_body(deployment, {**_NO_CACHE, "stream_options": stream_options}))
identity: Final = _assert_responses_restreamed(stream, iv.answer(deployment.name))
_assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
_assert_response_landed_once(deployment.model, identity)
def test_repeated_identical_streams_each_land_their_own_row(rig: _Rig, scenario: Scenario) -> None:
deployment: Final = _deploy(rig, scenario, "search")
body: Final = _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)})
identities: Final = frozenset(
_assert_restreamed(_chat(rig, body), iv.answer(deployment.name), usage=iv.FOLLOWUP_USAGE) for _ in range(3)
)
assert len(identities) == 3, identities
_assert_upstream(_upstream(rig, deployment), model_calls=6, searches=3)
_assert_landed_once(deployment.model, identities, rows=3)
@dataclass(frozen=True, slots=True)
class _Served:
path: str
stream: _Stream
async def _send(client: httpx.AsyncClient, key: str, path: str, body: Mapping[str, JsonValue]) -> _Served:
response: Final = await client.post(path, json=body, headers={"Authorization": f"Bearer {key}"})
return _Served(path, _Stream(response.status_code, response.text))
async def _burst(
base_url: str, key: str, requests: Sequence[tuple[str, Mapping[str, JsonValue]]], *, tolerate: bool = False
) -> tuple[_Served, ...]:
async with httpx.AsyncClient(base_url=base_url, timeout=90, trust_env=False) as client:
results: Final = await asyncio.gather(
*(_send(client, key, path, body) for path, body in requests), return_exceptions=tolerate
)
assert [result for result in results if isinstance(result, BaseException)] == [] or tolerate, results
assert all(isinstance(result, _Served | httpx.TransportError) for result in results), results
return tuple(result for result in results if isinstance(result, _Served))
def _chat_request(deployment: _Deployment) -> tuple[str, Mapping[str, JsonValue]]:
return "/v1/chat/completions", _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)})
def _responses_request(deployment: _Deployment) -> tuple[str, Mapping[str, JsonValue]]:
return "/v1/responses", _responses_body(deployment, {**_NO_CACHE, "stream_options": dict(_OBFUSCATION)})
async def test_concurrent_burst_across_chat_and_responses_answers_each_and_lands_each_once(
rig: _Rig, scenario: Scenario
) -> None:
chat: Final = _deploy(rig, scenario, "search")
responses: Final = _deploy(rig, scenario, "search")
outage: Final = _deploy(rig, scenario, "search-down")
requests: Final = (
*(_chat_request(chat) for _ in range(8)),
*(_responses_request(responses) for _ in range(8)),
*(_chat_request(outage) for _ in range(4)),
)
served: Final = await _burst(rig.base_url, rig.gateway.key, requests)
chat_ids: Final = frozenset(
_assert_restreamed(item.stream, iv.answer(chat.name), usage=iv.FOLLOWUP_USAGE) for item in served[:8]
)
response_ids: Final = tuple(
_assert_responses_restreamed(item.stream, iv.answer(responses.name)) for item in served[8:16]
)
outage_ids: Final = frozenset(
_assert_restreamed(item.stream, iv.answer(outage.name), usage=iv.FOLLOWUP_USAGE) for item in served[16:]
)
drained: Final = rig.wire.drain()
_assert_upstream(iv.received(drained, chat.name), model_calls=16, searches=8)
_assert_upstream(iv.received(drained, responses.name), model_calls=16, searches=8)
_assert_upstream(iv.received(drained, outage.name), model_calls=8, searches=4)
_assert_landed_once(chat.model, chat_ids, rows=8)
_assert_landed_once(outage.model, outage_ids, rows=4)
rows: Final = _spend_rows(responses.model, 8)
assert len(rows) == 8 and len(set(response_ids)) == 8, rows
assert all(
any(rv.same_response(string_value(row["request_id"]), identity) for row in rows) for identity in response_ids
)
def test_always_include_stream_usage_restreams_usage_on_chat_and_strips_it_from_azure_on_responses(
tmp_path: Path,
) -> None:
with _interception_proxy(tmp_path, iv.respond, {"always_include_stream_usage": True}) as owned:
with owned.gateway.scenario() as opened:
chat: Final = _deploy(owned, opened, "search")
responses: Final = _deploy(owned, opened, "search")
stream: Final = _chat(owned, _chat_body(chat))
identity: Final = _assert_restreamed(stream, iv.answer(chat.name), usage=iv.FOLLOWUP_USAGE)
replayed: Final = _assert_responses_restreamed(
_responses(owned, _responses_body(responses)), iv.answer(responses.name)
)
drained: Final = owned.wire.drain()
_assert_upstream(iv.received(drained, chat.name), model_calls=2, searches=1)
_assert_upstream(iv.received(drained, responses.name), model_calls=2, searches=1)
_assert_landed_once(chat.model, frozenset({identity}), rows=1)
_assert_response_landed_once(responses.model, replayed)
def _holding(release: threading.Event, held: SimpleQueue[str]) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
if request.method == "POST" and urlsplit(request.target).path.endswith("/chat/completions"):
held.put(request.target)
assert release.wait(timeout=90), "The burst was never released"
return iv.respond(request)
return respond
@contextmanager
def _releasing(release: threading.Event) -> Generator[None, None, None]:
try:
yield
finally:
release.set()
def _open_upstream_connections(pid: int, upstream: str) -> int:
port: Final = urlsplit(upstream).port
return sum(
1
for connection in psutil.Process(pid).net_connections(kind="tcp")
if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
)
def _worker_pids(log: str) -> tuple[int, ...]:
return tuple(
int(line.split(_STARTED_WORKER, 1)[1].split("]", 1)[0]) for line in log.splitlines() if _STARTED_WORKER in line
)
def _burst_models(wire: Wire) -> Sequence[Mapping[str, JsonValue]]:
return (
{
"model_name": _BURST_MODEL,
"litellm_params": {
"model": f"azure/{_BURST_DEPLOYMENT}",
"api_base": wire.url,
"api_key": iv.AZURE_KEY,
"api_version": iv.API_VERSION,
},
},
)
async def test_worker_sigkill_mid_burst_leaves_the_sibling_restreaming_with_usage(tmp_path: Path) -> None:
release: Final = threading.Event()
held: Final[SimpleQueue[str]] = SimpleQueue()
deployment: Final = _Deployment(_BURST_MODEL, _BURST_DEPLOYMENT)
with _interception_proxy(tmp_path, _holding(release, held), {}, _burst_models) as owned, _releasing(release):
workers: Final = eventually(
lambda: _worker_pids(owned.proxy.log.read_text()),
lambda pids: len(pids) == 2,
seconds=30,
)
burst: Final = asyncio.create_task(
_burst(
owned.base_url, owned.gateway.key, tuple(_chat_request(deployment) for _ in range(20)), tolerate=True
)
)
await asyncio.to_thread(eventually, held.qsize, lambda size: size == 20, 60)
held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, owned.wire.url) for pid in workers})
assert sum(held_by.values()) == 20, held_by
victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
victim: Final = psutil.Process(victim_pid)
victim.suspend()
victim.send_signal(signal.SIGKILL)
release.set()
served: Final = await burst
assert held_by[survivor_pid] >= 10, held_by
assert len(served) == held_by[survivor_pid], (held_by, len(served))
identities: Final = frozenset(
_assert_restreamed(item.stream, iv.answer(_BURST_DEPLOYMENT), usage=iv.FOLLOWUP_USAGE) for item in served
)
(follow_up,) = await _burst(owned.base_url, owned.gateway.key, (_chat_request(deployment),))
recovered: Final = _assert_restreamed(follow_up.stream, iv.answer(_BURST_DEPLOYMENT), usage=iv.FOLLOWUP_USAGE)
rows: Final = _spend_rows(_BURST_MODEL, len(identities) + 1)
landed: Final = tuple(string_value(row["request_id"]) for row in rows if row["status"] == "success")
assert sorted(landed) == sorted((*identities, recovered)), rows
replaced: Final = eventually(
lambda: owned.proxy.log.read_text(),
lambda log: len(_worker_pids(log)) == 3 and log.count(_STARTUP_COMPLETE) == 3,
seconds=2 * graceful_stop_seconds(),
)
assert victim_pid not in _worker_pids(replaced)[2:], replaced

97
tests/integration/translation/decisions/bases/cloudflare.py
View file

@ -6,6 +6,103 @@ from integration.translation.case import TranslationTestCase
"""
CLEF_TEST_CASE: Final = TranslationTestCase(
scenario="basic",
litellm_endpoint="/v1/decisions",
litellm_request={
"model": "cloudflare/@cf/cloudflare/clef",
"input": "Ticket (billing): The export job hangs at 99% and never finishes",
"questions": [
{"type": "predicate", "name": "defect", "instructions": "Is this a defect?"},
{
"type": "choice",
"name": "severity",
"instructions": "How severe is it?",
"choices": [
{"value": "low", "description": "cosmetic"},
{"value": "high", "description": "blocks users"},
],
},
{
"type": "score",
"name": "confidence",
"instructions": "How sure are you?",
"levels": [{"label": "unsure"}, {"label": "sure"}],
},
],
"cache": {"no-cache": True},
},
expected_provider_endpoint="/ai/run/@cf/cloudflare/clef",
expected_provider_headers={"authorization": "Bearer synthetic-cloudflare-key", "content-type": "application/json"},
expected_provider_request={
"model": "clef",
"state": "Ticket (billing): The export job hangs at 99% and never finishes",
"questions": {
"defect": {"type": "noul", "instructions": "Is this a defect?"},
"severity": {
"type": "choice",
"instructions": "How severe is it?",
"criteria": {"low": "cosmetic", "high": "blocks users"},
},
"confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]},
},
},
mock_provider_response={
"result": {
"model": "clef",
"answers": {
"defect": {"type": "noul", "noul": 0.9345},
"severity": {
"type": "choice",
"choice": "high",
"confidence": 0.8067,
"probabilities": {"low": 0.0509, "high": 0.9491},
},
"confidence": {
"type": "score",
"score": 0.9036,
"confidence": 0.6515,
"legend": {"0": "unsure", "1": "sure"},
"probabilities": {"0": 0.0964, "1": 0.9036},
},
},
"usage": {"input_tokens": 290, "output_tokens": 0},
},
"success": True,
"errors": [],
"messages": [],
},
expected_litellm_response={
"model": "clef",
"answers": [
{"type": "predicate", "name": "defect", "probability": 0.9345},
{
"type": "choice",
"name": "severity",
"choice": "high",
"probabilities": [{"value": "low", "probability": 0.0509}, {"value": "high", "probability": 0.9491}],
"confidence": 0.8067,
},
{
"type": "score",
"name": "confidence",
"score": 0.9036,
"probabilities": [
{"value": 0, "label": "unsure", "probability": 0.0964},
{"value": 1, "label": "sure", "probability": 0.9036},
],
"confidence": 0.6515,
},
],
"usage": {
"input_tokens": 290,
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
"output_tokens": 0,
"output_tokens_details": {"reasoning_tokens": 0},
"total_tokens": 290,
},
},
)
CLEF_SYSTEMONE_TEST_CASE: Final = TranslationTestCase(
scenario="systemone",
litellm_endpoint="/v1/systemone",
litellm_request={
"model": "cloudflare/@cf/cloudflare/clef",

94
tests/integration/translation/decisions/bases/openrouter.py
View file

@ -6,6 +6,100 @@ from integration.translation.case import TranslationTestCase
"""
TYPESAFE_JEV_1_13_TEST_CASE: Final = TranslationTestCase(
scenario="basic",
litellm_endpoint="/v1/decisions",
litellm_request={
"model": "openrouter/typesafe/jev-1.13",
"input": "Ticket (billing): The export job hangs at 99% and never finishes",
"questions": [
{"type": "predicate", "name": "defect", "instructions": "Is this a defect?"},
{
"type": "choice",
"name": "severity",
"instructions": "How severe is it?",
"choices": [
{"value": "low", "description": "cosmetic"},
{"value": "high", "description": "blocks users"},
],
},
{
"type": "score",
"name": "confidence",
"instructions": "How sure are you?",
"levels": [{"label": "unsure"}, {"label": "sure"}],
},
],
"cache": {"no-cache": True},
},
expected_provider_endpoint="/alpha/decisions",
expected_provider_headers={"authorization": "Bearer synthetic-openrouter-key", "content-type": "application/json"},
expected_provider_request={
"model": "typesafe/jev-1.13",
"state": "Ticket (billing): The export job hangs at 99% and never finishes",
"questions": {
"defect": {"type": "noul", "instructions": "Is this a defect?"},
"severity": {
"type": "choice",
"instructions": "How severe is it?",
"criteria": {"low": "cosmetic", "high": "blocks users"},
},
"confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]},
},
},
mock_provider_response={
"model": "typesafe/jev-1.13-20260917",
"answers": {
"defect": {"type": "noul", "noul": 0.81},
"severity": {
"type": "choice",
"choice": "high",
"confidence": 0.99,
"probabilities": {"low": 0.01, "high": 0.99},
},
"confidence": {
"type": "score",
"score": 0.5,
"confidence": 0,
"legend": {"0": "unsure", "1": "sure"},
"probabilities": {"0": 0.5, "1": 0.5},
},
},
"usage": {"input_tokens": 377, "output_tokens": 62, "cost": 1.5834e-05},
"id": "gen-dec-1791323839-GA15kY0nt34oiJ7srfki",
"provider": "TypeSafe",
},
expected_litellm_response={
"model": "typesafe/jev-1.13-20260917",
"answers": [
{"type": "predicate", "name": "defect", "probability": 0.81},
{
"type": "choice",
"name": "severity",
"choice": "high",
"probabilities": [{"value": "low", "probability": 0.01}, {"value": "high", "probability": 0.99}],
"confidence": 0.99,
},
{
"type": "score",
"name": "confidence",
"score": 0.5,
"probabilities": [
{"value": 0, "label": "unsure", "probability": 0.5},
{"value": 1, "label": "sure", "probability": 0.5},
],
"confidence": 0,
},
],
"usage": {
"input_tokens": 377,
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
"output_tokens": 62,
"output_tokens_details": {"reasoning_tokens": 0},
"total_tokens": 439,
},
},
)
TYPESAFE_JEV_1_13_SYSTEMONE_TEST_CASE: Final = TranslationTestCase(
scenario="systemone",
litellm_endpoint="/v1/systemone",
litellm_request={
"model": "openrouter/typesafe/jev-1.13",

95
tests/integration/translation/decisions/bases/perplexity.py
View file

@ -6,6 +6,101 @@ from integration.translation.case import TranslationTestCase
"""
PPLX_DECIDER_V1_27B_TEST_CASE: Final = TranslationTestCase(
scenario="basic",
litellm_endpoint="/v1/decisions",
litellm_request={
"model": "perplexity/pplx-decider-v1-27b",
"input": "Ticket (billing): The export job hangs at 99% and never finishes",
"questions": [
{"type": "predicate", "name": "defect", "instructions": "Is this a defect?"},
{
"type": "choice",
"name": "severity",
"instructions": "How severe is it?",
"choices": [
{"value": "low", "description": "cosmetic"},
{"value": "high", "description": "blocks users"},
],
},
{
"type": "score",
"name": "confidence",
"instructions": "How sure are you?",
"levels": [{"label": "unsure"}, {"label": "sure"}],
},
],
"cache": {"no-cache": True},
},
expected_provider_endpoint="/v1/decisions",
expected_provider_headers={"authorization": "Bearer synthetic-perplexity-key", "content-type": "application/json"},
expected_provider_request={
"model": "pplx-decider-v1-27b",
"state": "Ticket (billing): The export job hangs at 99% and never finishes",
"questions": {
"defect": {"type": "noul", "instructions": "Is this a defect?"},
"severity": {
"type": "choice",
"instructions": "How severe is it?",
"criteria": {"low": "cosmetic", "high": "blocks users"},
},
"confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]},
},
},
mock_provider_response={
"model": "pplx-decider-v1-27b",
"answers": {
"defect": {"type": "noul", "noul": 0.9989100737587077},
"severity": {
"type": "choice",
"choice": "high",
"confidence": 0.9964631215356778,
"probabilities": {"low": 0.0017684392321610232, "high": 0.9982315607678389},
},
"confidence": {
"type": "score",
"score": 0.07367392327139817,
"confidence": 0.8526521534572037,
"legend": {"0": "unsure", "1": "sure"},
"probabilities": {"0": 0.9263260767286018, "1": 0.07367392327139817},
},
},
"usage": {"input_tokens": 318, "output_tokens": 3},
},
expected_litellm_response={
"model": "pplx-decider-v1-27b",
"answers": [
{"type": "predicate", "name": "defect", "probability": 0.9989100737587077},
{
"type": "choice",
"name": "severity",
"choice": "high",
"probabilities": [
{"value": "low", "probability": 0.0017684392321610232},
{"value": "high", "probability": 0.9982315607678389},
],
"confidence": 0.9964631215356778,
},
{
"type": "score",
"name": "confidence",
"score": 0.07367392327139817,
"probabilities": [
{"value": 0, "label": "unsure", "probability": 0.9263260767286018},
{"value": 1, "label": "sure", "probability": 0.07367392327139817},
],
"confidence": 0.8526521534572037,
},
],
"usage": {
"input_tokens": 318,
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
"output_tokens": 3,
"output_tokens_details": {"reasoning_tokens": 0},
"total_tokens": 321,
},
},
)
PPLX_DECIDER_V1_27B_SYSTEMONE_TEST_CASE: Final = TranslationTestCase(
scenario="systemone",
litellm_endpoint="/v1/systemone",
litellm_request={
"model": "perplexity/pplx-decider-v1-27b",

34
tests/integration/translation/decisions/bases/strands_decider.py
View file

@ -6,6 +6,40 @@ from integration.translation.case import TranslationTestCase
"""
STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE: Final = TranslationTestCase(
scenario="basic",
litellm_endpoint="/v1/decisions",
litellm_request={
"model": "strands_decider/strands-decider-2B-hobson-v19",
"input": "Help! My payouts have been failing for 3 days!",
"questions": [{"type": "predicate", "name": "is_urgent", "instructions": "Does this convey urgency?"}],
"cache": {"no-cache": True},
},
expected_provider_endpoint="/v1/systemone",
expected_provider_headers={"content-type": "application/json"},
expected_provider_request={
"model": "strands-decider-2B-hobson-v19",
"state": "Help! My payouts have been failing for 3 days!",
"questions": {"is_urgent": {"type": "noul", "instructions": "Does this convey urgency?"}},
},
mock_provider_response={
"model": "strands-decider-2B-hobson-v19",
"answers": {"is_urgent": {"type": "noul", "noul": 0.8277}},
"usage": {"input_tokens": 86, "output_tokens": 1},
"latency_ms": 140.03,
},
expected_litellm_response={
"model": "strands-decider-2B-hobson-v19",
"answers": [{"type": "predicate", "name": "is_urgent", "probability": 0.8277}],
"usage": {
"input_tokens": 86,
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
"output_tokens": 1,
"output_tokens_details": {"reasoning_tokens": 0},
"total_tokens": 87,
},
},
)
STRANDS_DECIDER_2B_HOBSON_V19_SYSTEMONE_TEST_CASE: Final = TranslationTestCase(
scenario="systemone",
litellm_endpoint="/v1/systemone",
litellm_request={
"model": "strands_decider/strands-decider-2B-hobson-v19",

92
tests/integration/translation/decisions/bases/typesafe.py
View file

@ -6,6 +6,98 @@ from integration.translation.case import TranslationTestCase
"""
JEV_1_13_0_TEST_CASE: Final = TranslationTestCase(
scenario="basic",
litellm_endpoint="/v1/decisions",
litellm_request={
"model": "typesafe/jev-1.13.0",
"input": "Ticket (billing): The export job hangs at 99% and never finishes",
"questions": [
{"type": "predicate", "name": "defect", "instructions": "Is this a defect?"},
{
"type": "choice",
"name": "severity",
"instructions": "How severe is it?",
"choices": [
{"value": "low", "description": "cosmetic"},
{"value": "high", "description": "blocks users"},
],
},
{
"type": "score",
"name": "confidence",
"instructions": "How sure are you?",
"levels": [{"label": "unsure"}, {"label": "sure"}],
},
],
"cache": {"no-cache": True},
},
expected_provider_endpoint="/v1/systemone",
expected_provider_headers={"authorization": "Bearer synthetic-typesafe-key", "content-type": "application/json"},
expected_provider_request={
"model": "jev-1.13.0",
"state": "Ticket (billing): The export job hangs at 99% and never finishes",
"questions": {
"defect": {"type": "noul", "instructions": "Is this a defect?"},
"severity": {
"type": "choice",
"instructions": "How severe is it?",
"criteria": {"low": "cosmetic", "high": "blocks users"},
},
"confidence": {"type": "score", "instructions": "How sure are you?", "criteria": ["unsure", "sure"]},
},
},
mock_provider_response={
"model": "jev-1.13.0",
"answers": {
"defect": {"type": "noul", "noul": 0.78},
"severity": {
"type": "choice",
"choice": "high",
"confidence": 0.99,
"probabilities": {"low": 0.01, "high": 0.99},
},
"confidence": {
"type": "score",
"score": 0.51,
"confidence": 0.03,
"legend": {"0": "unsure", "1": "sure"},
"probabilities": {"0": 0.49, "1": 0.51},
},
},
"usage": {"input_tokens": 377, "output_tokens": 62},
},
expected_litellm_response={
"model": "jev-1.13.0",
"answers": [
{"type": "predicate", "name": "defect", "probability": 0.78},
{
"type": "choice",
"name": "severity",
"choice": "high",
"probabilities": [{"value": "low", "probability": 0.01}, {"value": "high", "probability": 0.99}],
"confidence": 0.99,
},
{
"type": "score",
"name": "confidence",
"score": 0.51,
"probabilities": [
{"value": 0, "label": "unsure", "probability": 0.49},
{"value": 1, "label": "sure", "probability": 0.51},
],
"confidence": 0.03,
},
],
"usage": {
"input_tokens": 377,
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
"output_tokens": 62,
"output_tokens_details": {"reasoning_tokens": 0},
"total_tokens": 439,
},
},
)
JEV_1_13_0_SYSTEMONE_TEST_CASE: Final = TranslationTestCase(
scenario="systemone",
litellm_endpoint="/v1/systemone",
litellm_request={
"model": "typesafe/jev-1.13.0",

4
tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py
View file

@ -2,10 +2,10 @@ import pytest
from integration._support.client import Gateway
from integration._support.provider import SharedProvider
from integration.translation.case import TranslationTestCase
from integration.translation.decisions.bases.cloudflare import CLEF_TEST_CASE
from integration.translation.decisions.bases.cloudflare import CLEF_TEST_CASE, CLEF_SYSTEMONE_TEST_CASE
from integration.translation.runner import assert_translation
@pytest.mark.parametrize("case", [CLEF_TEST_CASE], ids=lambda case: case.id)
@pytest.mark.parametrize("case", [CLEF_TEST_CASE, CLEF_SYSTEMONE_TEST_CASE], ids=lambda case: case.id)
def test_decisions_basic_cloudflare(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None:
assert_translation(case, gateway, provider)

9
tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py
View file

@ -2,10 +2,15 @@ import pytest
from integration._support.client import Gateway
from integration._support.provider import SharedProvider
from integration.translation.case import TranslationTestCase
from integration.translation.decisions.bases.openrouter import TYPESAFE_JEV_1_13_TEST_CASE
from integration.translation.decisions.bases.openrouter import (
TYPESAFE_JEV_1_13_TEST_CASE,
TYPESAFE_JEV_1_13_SYSTEMONE_TEST_CASE,
)
from integration.translation.runner import assert_translation
@pytest.mark.parametrize("case", [TYPESAFE_JEV_1_13_TEST_CASE], ids=lambda case: case.id)
@pytest.mark.parametrize(
"case", [TYPESAFE_JEV_1_13_TEST_CASE, TYPESAFE_JEV_1_13_SYSTEMONE_TEST_CASE], ids=lambda case: case.id
)
def test_decisions_basic_openrouter(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None:
assert_translation(case, gateway, provider)

9
tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py
View file

@ -2,10 +2,15 @@ import pytest
from integration._support.client import Gateway
from integration._support.provider import SharedProvider
from integration.translation.case import TranslationTestCase
from integration.translation.decisions.bases.perplexity import PPLX_DECIDER_V1_27B_TEST_CASE
from integration.translation.decisions.bases.perplexity import (
PPLX_DECIDER_V1_27B_TEST_CASE,
PPLX_DECIDER_V1_27B_SYSTEMONE_TEST_CASE,
)
from integration.translation.runner import assert_translation
@pytest.mark.parametrize("case", [PPLX_DECIDER_V1_27B_TEST_CASE], ids=lambda case: case.id)
@pytest.mark.parametrize(
"case", [PPLX_DECIDER_V1_27B_TEST_CASE, PPLX_DECIDER_V1_27B_SYSTEMONE_TEST_CASE], ids=lambda case: case.id
)
def test_decisions_basic_perplexity(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None:
assert_translation(case, gateway, provider)

11
tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py
View file

@ -2,10 +2,17 @@ import pytest
from integration._support.client import Gateway
from integration._support.provider import SharedProvider
from integration.translation.case import TranslationTestCase
from integration.translation.decisions.bases.strands_decider import STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE
from integration.translation.decisions.bases.strands_decider import (
STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE,
STRANDS_DECIDER_2B_HOBSON_V19_SYSTEMONE_TEST_CASE,
)
from integration.translation.runner import assert_translation
@pytest.mark.parametrize("case", [STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE], ids=lambda case: case.id)
@pytest.mark.parametrize(
"case",
[STRANDS_DECIDER_2B_HOBSON_V19_TEST_CASE, STRANDS_DECIDER_2B_HOBSON_V19_SYSTEMONE_TEST_CASE],
ids=lambda case: case.id,
)
def test_decisions_basic_strands_decider(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None:
assert_translation(case, gateway, provider)

4
tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py
View file

@ -2,10 +2,10 @@ import pytest
from integration._support.client import Gateway
from integration._support.provider import SharedProvider
from integration.translation.case import TranslationTestCase
from integration.translation.decisions.bases.typesafe import JEV_1_13_0_TEST_CASE
from integration.translation.decisions.bases.typesafe import JEV_1_13_0_TEST_CASE, JEV_1_13_0_SYSTEMONE_TEST_CASE
from integration.translation.runner import assert_translation
@pytest.mark.parametrize("case", [JEV_1_13_0_TEST_CASE], ids=lambda case: case.id)
@pytest.mark.parametrize("case", [JEV_1_13_0_TEST_CASE, JEV_1_13_0_SYSTEMONE_TEST_CASE], ids=lambda case: case.id)
def test_decisions_basic_typesafe(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None:
assert_translation(case, gateway, provider)

32
tests/test_litellm_rust/messages/test_callbacks.py
View file

@ -239,3 +239,35 @@ async def test_native_messages_stream_success_log_carries_usage_rebuilt_from_the
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"
@pytest.mark.asyncio
async def test_native_messages_does_not_restore_keywords_deleted_by_a_pre_call_hook(
messages_server: RecordingServer,
) -> None:
class DropKeywords(CustomLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: object
) -> dict[str, object]:
return {name: value for name, value in kwargs.items() if name not in ("system", "extra_headers")}
call_arguments: Final = arguments(
messages_server,
system="remove this instruction",
extra_headers={"x-deleted": "remove this header"},
headers={"x-kept": "keep this header"},
stream=False,
)
with rebound(litellm, "callbacks", [DropKeywords()]):
await litellm.anthropic.messages.acreate(**call_arguments)
assert_served_natively(messages_server)
sent: Final = messages_server.requests[0]
assert sent.body == {
"model": MESSAGES_MODEL.removeprefix("anthropic/"),
"messages": call_arguments["messages"],
"max_tokens": call_arguments["max_tokens"],
"stream": call_arguments["stream"],
}
assert sent.headers.get("x-deleted") is None
assert sent.headers.get("x-kept") == "keep this header"

19
tests/test_litellm_rust/ocr/test_callbacks.py
View file

@ -16,7 +16,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
from tests.test_litellm_rust.support.isolation import isolated_callback_registries
from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
from tests.test_litellm_rust.support.requests import (
OCR_DOCUMENT,
@ -65,6 +65,23 @@ def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_
}
@pytest.mark.asyncio
async def test_native_ocr_does_not_restore_keywords_deleted_by_a_pre_call_hook(ocr_server: RecordingServer) -> None:
class DropKeywords(CustomLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: object
) -> dict[str, object]:
return {name: value for name, value in kwargs.items() if name not in ("pages", "extra_headers")}
with rebound(litellm, "callbacks", [DropKeywords()]):
await call_native_aocr(ocr_server, pages=[0], extra_headers={"x-deleted": "remove this header"})
assert len(ocr_server.requests) == 1
sent: Final = ocr_server.requests[0]
assert sent.body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT}
assert sent.headers.get("x-deleted") is None
@pytest.mark.parametrize("raise_after_edit", [False, True], ids=["callback-returns", "callback-raises"])
def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider(
ocr_server: RecordingServer, raise_after_edit: bool

22
tests/test_litellm_rust/test_inference.py
View file

@ -87,7 +87,6 @@ def native_call(
"api_base": server.base_url,
"custom_llm_provider": "openai",
"extra_headers": None,
**response_kwargs,
},
)
return (_native.aresponses if asynchronous else _native.responses)(response_request)
@ -147,6 +146,27 @@ async def test_native_inference_pre_call_edits_reach_the_provider(
assert _OBJECT.validate_python(recording_server.requests[0].body)["temperature"] == 0.75
@pytest.mark.asyncio
async def test_native_inference_does_not_restore_keywords_deleted_by_a_pre_call_hook(
route: Route, recording_server: RecordingServer
) -> None:
class DropKeywords(CustomLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: CallTypes | None
) -> dict[str, object]:
return {name: value for name, value in kwargs.items() if name not in ("temperature", "extra_headers")}
litellm.callbacks.append(DropKeywords())
await execute(
route, True, recording_server, {"temperature": 0.25, "extra_headers": {"x-deleted": "remove this header"}}
)
assert len(recording_server.requests) == 1
sent: Final = recording_server.requests[0]
assert "temperature" not in _OBJECT.validate_python(sent.body)
assert sent.headers.get("x-deleted") is None
@pytest.mark.asyncio
@pytest.mark.parametrize("from_credentials", (False, True))
async def test_native_resource_setup_uses_deployment_hook_arguments(

66
tests/typing/openai_sdk_compat.py Normal file
View file

@ -0,0 +1,66 @@
from typing import Final
import httpx
from openai import DefaultAsyncHttpxClient, DefaultHttpxClient
from openai import HttpxBinaryResponseContent as SDKBinaryResponse
from openai import Timeout as SDKTimeout
from openai._types import Response as SDKResponse
from typing_extensions import assert_type
import litellm
from litellm.exceptions import (
APIConnectionError,
APIResponseValidationError,
BadGatewayError,
InternalServerError,
InvalidRequestError,
RateLimitError,
Timeout,
)
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
from litellm.types.llms.openai import HttpxBinaryResponseContent
from litellm.types.router import GenericLiteLLMParams, LiteLLMParamsTypedDict
def consume_binary_response(response: HttpxBinaryResponseContent) -> httpx.Response:
assert_type(response.response, httpx.Response)
assert_type(response.response.request, httpx.Request)
return response.response
def binary_response_types(legacy_response: httpx.Response, sdk_response: SDKResponse) -> None:
legacy: Final = HttpxBinaryResponseContent(legacy_response)
native: Final = HttpxBinaryResponseContent(sdk_response)
assert_type(legacy.response, httpx.Response)
assert_type(native.response, SDKResponse)
assert_type(legacy.read(), bytes)
assert_type(native.read(), bytes)
assert_type(consume_binary_response(legacy), httpx.Response)
sdk_wrapper: Final[SDKBinaryResponse] = legacy
assert_type(sdk_wrapper.read(), bytes)
def synthesized_response_types(
error: RateLimitError | BadGatewayError | InternalServerError | APIResponseValidationError | InvalidRequestError,
) -> None:
assert_type(error.response, httpx.Response)
assert_type(error.request, httpx.Request)
def synthesized_request_types(error: APIConnectionError | Timeout) -> None:
assert_type(error.request, httpx.Request)
def configured_clients(
sync_client: httpx.Client | DefaultHttpxClient,
async_client: httpx.AsyncClient | DefaultAsyncHttpxClient,
timeout: httpx.Timeout | SDKTimeout,
) -> None:
litellm.client_session = sync_client # test-quality-ok: [TQ005] Type-check-only fixture; never executed
litellm.aclient_session = async_client # test-quality-ok: [TQ005] Type-check-only fixture; never executed
assert_type(CompletionTimeout.normalize(timeout), httpx.Timeout)
params: Final[LiteLLMParamsTypedDict] = {"timeout": 1.0}
params["timeout"] = timeout
GenericLiteLLMParams(timeout=timeout)

8
tests/typing/pyrightconfig.json Normal file
View file

@ -0,0 +1,8 @@
{
"extends": "../../pyrightconfig.json",
"include": ["openai_sdk_compat.py"],
"exclude": [],
"extraPaths": ["../.."],
"pythonVersion": "3.10",
"reportMissingImports": true
}

335
tests/unit/decisions/test_main.py
View file

@ -2,6 +2,8 @@ from __future__ import annotations
import asyncio
import json
import threading
from concurrent.futures import wait
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
@ -9,11 +11,16 @@ from typing import Final
import httpx
import pytest
import respx
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import Credentials
import litellm
from litellm.cost_calculator import get_response_cost_from_hidden_params
from litellm.constants import AWS_SIGNING_MAX_THREADS
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.llms.bedrock.base_aws_llm import AWS_SIGNING_EXECUTOR
from litellm.types.decisions import (
ChoiceAnswer,
DecisionsResponse,
@ -698,6 +705,33 @@ async def test_strands_decider_without_key_preserves_response_extras(
assert severity.legend == {"0": "none", "1": "low", "2": "high"}
@pytest.mark.asyncio
async def test_strands_decider_plain_api_base_answers_while_every_aws_signing_thread_is_busy(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False)
respx_mock.post("https://strands.example/v1/systemone").respond(json=_STRANDS_RESPONSE)
release: Final = threading.Event()
busy: Final = tuple(AWS_SIGNING_EXECUTOR.submit(release.wait) for _ in range(AWS_SIGNING_MAX_THREADS))
try:
response: Final = await asyncio.wait_for(
litellm.adecisions(
model="strands_decider/strands-decider-2B-hobson-v19",
state="review",
questions={"severity": {"type": "score", "criteria": ["none", "low", "high"]}},
api_base="https://strands.example",
),
timeout=5,
)
finally:
release.set()
wait(busy)
assert isinstance(response.answers["severity"], ScoreAnswer)
@pytest.mark.asyncio
async def test_strands_decider_uses_key_from_matching_environment_base(
monkeypatch: pytest.MonkeyPatch,
@ -750,6 +784,172 @@ async def test_strands_decider_provider_resolution_and_router_dispatch(
assert response.model == _STRANDS_RESPONSE["model"]
_RUNTIME_ARN: Final = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/qa_decider-AbC123xyz0"
_RUNTIME_URL: Final = (
"https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/"
"arn%3Aaws%3Abedrock-agentcore%3Aus-east-1%3A123456789012%3Aruntime%2Fqa_decider-AbC123xyz0/invocations"
)
_RUNTIME_CREDENTIALS: Final = Credentials(access_key="AKIDEXAMPLE", secret_key="example-secret")
_SESSION_HEADER: Final = "x-amzn-bedrock-agentcore-runtime-session-id"
def _sigv4_verifies(request: httpx.Request) -> bool:
authorization: Final = request.headers["authorization"]
signed_names: Final = authorization.split("SignedHeaders=")[1].split(",")[0].split(";")
replayed: Final = AWSRequest(
method=request.method,
url=str(request.url),
data=request.content,
headers={name: request.headers[name] for name in signed_names},
)
replayed.context["timestamp"] = request.headers["x-amz-date"]
auth: Final = SigV4Auth(_RUNTIME_CREDENTIALS, "bedrock-agentcore", "us-east-1")
string_to_sign: Final = auth.string_to_sign(replayed, auth.canonical_request(replayed))
return authorization.endswith(f"Signature={auth.signature(string_to_sign, replayed)}")
@pytest.mark.parametrize("use_async", [True, False])
@pytest.mark.asyncio
async def test_strands_decider_runtime_arn_sends_a_sigv4_signed_invoke_agent_runtime_call(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
use_async: bool,
) -> None:
monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False)
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
route: Final = respx_mock.post(_RUNTIME_URL).respond(json=_STRANDS_RESPONSE)
call_args: Final[Mapping[str, object]] = {
"model": "strands_decider/strands-decider-2B-hobson-v19",
"state": "review",
"questions": {"severity": {"type": "score", "criteria": ["none", "low", "high"]}},
"api_base": _RUNTIME_ARN,
"aws_access_key_id": _RUNTIME_CREDENTIALS.access_key,
"aws_secret_access_key": _RUNTIME_CREDENTIALS.secret_key,
}
response: Final = await litellm.adecisions(**call_args) if use_async else litellm.decisions(**call_args)
assert route.call_count == 1
sent: Final = route.calls[0].request
assert _sigv4_verifies(sent)
assert json.loads(sent.content)["model"] == "strands-decider-2B-hobson-v19"
assert sent.headers[_SESSION_HEADER].startswith("litellm-decider-")
assert response.model == _STRANDS_RESPONSE["model"]
@pytest.mark.asyncio
async def test_strands_decider_runtime_arn_router_deployment_answers_openai_format_calls(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False)
router: Final = litellm.Router(
model_list=[
{
"model_name": "strands",
"litellm_params": {
"model": "strands_decider/strands-decider-2B-hobson-v19",
"api_base": _RUNTIME_ARN,
"aws_access_key_id": _RUNTIME_CREDENTIALS.access_key,
"aws_secret_access_key": _RUNTIME_CREDENTIALS.secret_key,
"aws_region_name": "eu-west-1",
"agentcore_runtime_session_id": "pooled-session-0123456789abcdef0123456",
},
}
]
)
route: Final = respx_mock.post(_RUNTIME_URL).respond(json=_STRANDS_RESPONSE)
response: Final = await router.adecisions(
model="strands",
input="review",
extra_headers={_SESSION_HEADER: "caller-session-0123456789abcdef01234567"},
questions=[
{
"type": "score",
"name": "severity",
"instructions": "How bad?",
"levels": [{"label": "none"}, {"label": "low"}, {"label": "high"}],
}
],
)
sent: Final = route.calls[0].request
assert _sigv4_verifies(sent)
assert sent.headers[_SESSION_HEADER] == "pooled-session-0123456789abcdef0123456"
assert isinstance(response, OpenAIDecisionResponse)
@pytest.mark.asyncio
async def test_strands_decider_runtime_arn_with_api_key_sends_the_bearer_token_unsigned(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False)
route: Final = respx_mock.post(_RUNTIME_URL).respond(json=_STRANDS_RESPONSE)
await litellm.adecisions(
model="strands_decider/strands-decider-2B-hobson-v19",
state="review",
questions={"severity": {"type": "score", "criteria": ["none", "low", "high"]}},
api_base=_RUNTIME_ARN,
api_key="runtime-jwt",
)
sent: Final = route.calls[0].request
assert sent.headers["authorization"] == "Bearer runtime-jwt"
assert "x-amz-date" not in sent.headers
assert _SESSION_HEADER in sent.headers
@pytest.mark.parametrize(
("code", "error_class"),
[
("bad_request", litellm.BadRequestError),
("loading", litellm.ServiceUnavailableError),
("inference_error", litellm.InternalServerError),
],
)
@pytest.mark.asyncio
async def test_strands_decider_runtime_error_envelope_raises_instead_of_answering(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
code: str,
error_class: type[Exception],
) -> None:
monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False)
respx_mock.post(_RUNTIME_URL).respond(json={"error": {"code": code, "message": "at most 16 questions"}})
with pytest.raises(error_class, match="at most 16 questions"):
await litellm.adecisions(
model="strands_decider/strands-decider-2B-hobson-v19",
state="review",
questions={"severity": {"type": "score", "criteria": ["none", "low", "high"]}},
api_base=_RUNTIME_ARN,
aws_access_key_id=_RUNTIME_CREDENTIALS.access_key,
aws_secret_access_key=_RUNTIME_CREDENTIALS.secret_key,
)
@pytest.mark.asyncio
async def test_strands_decider_plain_api_base_error_envelope_stays_an_unexpected_response(
monkeypatch: pytest.MonkeyPatch,
respx_mock: respx.MockRouter,
) -> None:
monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False)
respx_mock.post("https://strands.example/v1/systemone").respond(
json={"error": {"code": "bad_request", "message": "at most 16 questions"}}
)
with pytest.raises(litellm.InternalServerError, match="returned an unexpected response"):
await litellm.adecisions(
model="strands_decider/strands-decider-2B-hobson-v19",
state="review",
questions={"severity": {"type": "score", "criteria": ["none", "low", "high"]}},
api_base="https://strands.example",
)
_VLLM_RESPONSE: Final[Mapping[str, object]] = {
"id": "systemone-4af3d2c1",
"object": "structured_decision",
@ -1262,6 +1462,141 @@ async def test_openrouter_decisions_uses_provider_reported_cost_without_cost_map
assert get_response_cost_from_hidden_params(response.hidden_params) == cost
_SAFETY_IDENTIFIER: Final = "end-user-7"
_PREDICATE_QUESTIONS: Final[tuple[Mapping[str, object], ...]] = (
{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"},
)
_PREDICATE_REQUESTS: Final[tuple[Mapping[str, object], ...]] = (
MappingProxyType({"input": "review", "questions": _PREDICATE_QUESTIONS}),
MappingProxyType(
{"state": "review", "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}}
),
)
_PREDICATE_REQUEST_IDS: Final = ("input_format", "state_format")
@pytest.mark.parametrize("request_kwargs", _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS)
def test_safety_identifier_is_refused_before_http_when_the_provider_cannot_take_it(
request_kwargs: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "drop_params", False)
route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE)
with pytest.raises(litellm.UnsupportedParamsError, match=r"safety_identifier.*drop_params") as caught:
litellm.decisions(
model="perplexity/pplx-decider-v1-27b",
safety_identifier=_SAFETY_IDENTIFIER,
api_key="caller-key",
**request_kwargs,
)
assert caught.value.status_code == 400
assert not route.called
@pytest.mark.asyncio
@pytest.mark.parametrize("scope", ("call", "global"))
async def test_safety_identifier_is_dropped_from_the_wire_under_drop_params(
scope: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "drop_params", scope == "global")
route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE)
response: Final = await litellm.adecisions(
model="perplexity/pplx-decider-v1-27b",
input="review",
questions=_PREDICATE_QUESTIONS,
safety_identifier=_SAFETY_IDENTIFIER,
api_key="caller-key",
**({"drop_params": True} if scope == "call" else {}),
)
assert route.called
assert json.loads(respx_mock.calls[0].request.content) == {
"model": "pplx-decider-v1-27b",
"state": "review",
"questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}},
}
assert isinstance(response, OpenAIDecisionResponse)
assert [answer.model_dump(mode="json") for answer in response.answers] == [
{"type": "predicate", "name": "is_defect", "probability": 0.9}
]
@pytest.mark.asyncio
@pytest.mark.parametrize("request_kwargs", _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS)
@pytest.mark.parametrize("drop_params", (False, True), ids=("strict", "drop_params"))
async def test_safety_identifier_reaches_openai_whether_or_not_params_are_dropped(
drop_params: bool,
request_kwargs: Mapping[str, object],
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "drop_params", drop_params)
monkeypatch.setattr(litellm, "api_base", None)
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
monkeypatch.delenv("OPENAI_API_BASE", raising=False)
route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE)
await litellm.adecisions(
model="openai/gpt-6-luna",
safety_identifier=_SAFETY_IDENTIFIER,
api_key="caller-key",
**request_kwargs,
)
assert route.called
assert json.loads(respx_mock.calls[0].request.content) == {
"model": "gpt-6-luna",
"input": "review",
"questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}],
"safety_identifier": _SAFETY_IDENTIFIER,
}
@pytest.mark.asyncio
@pytest.mark.parametrize("request_kwargs", _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS)
async def test_non_string_safety_identifier_is_rejected_before_http(
request_kwargs: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "drop_params", False)
route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE)
invalid_safety_identifier: Final[Mapping[str, object]] = MappingProxyType({"safety_identifier": 7})
with pytest.raises(litellm.BadRequestError, match=r"(?s)Invalid Decisions request.*safety_identifier"):
await litellm.adecisions(
model="openai/gpt-6-luna", api_key="caller-key", **request_kwargs, **invalid_safety_identifier
)
assert not route.called
@pytest.mark.asyncio
@pytest.mark.parametrize("request_kwargs", _PREDICATE_REQUESTS, ids=_PREDICATE_REQUEST_IDS)
async def test_non_string_safety_identifier_is_dropped_from_the_wire_under_drop_params(
request_kwargs: Mapping[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "drop_params", False)
monkeypatch.setattr(litellm, "api_base", None)
monkeypatch.delenv("OPENAI_BASE_URL", raising=False)
monkeypatch.delenv("OPENAI_API_BASE", raising=False)
route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE)
dropped_safety_identifier: Final[Mapping[str, object]] = MappingProxyType(
{"safety_identifier": 7, "drop_params": True}
)
await litellm.adecisions(
model="openai/gpt-6-luna", api_key="caller-key", **request_kwargs, **dropped_safety_identifier
)
assert route.called
assert json.loads(respx_mock.calls[0].request.content) == {
"model": "gpt-6-luna",
"input": "review",
"questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}],
}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"api_base",

30
tests/unit/integrations/code_interpreter_interception/test_handler.py
View file

@ -6,22 +6,26 @@ config and records how it is called.
"""
import time
from typing import Final
import pytest
from litellm.integrations.code_interpreter_interception.handler import (
CodeInterpreterInterceptionLogger,
LITELLM_CODE_EXECUTION_TOOL_NAME,
_INTERCEPTION_ACTIVE_KEY as _ACTIVE_KEY,
)
from litellm.integrations.code_interpreter_interception.handler import (
_SANDBOX_KEY,
_SESSION_SCOPED_KEY,
LITELLM_CODE_EXECUTION_TOOL_NAME,
CodeInterpreterInterceptionLogger,
)
from litellm.llms.base_llm.sandbox.transformation import CodeExecutionResult
from litellm.types.integrations.custom_logger import (
CHAT_COMPLETION_AGENTIC_SURFACE,
CODE_INTERPRETER_STREAM_OPTIONS_KEY,
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
is_interception_internal_key,
)
from litellm.llms.base_llm.sandbox.transformation import CodeExecutionResult
from litellm.types.utils import CallTypes
@ -516,6 +520,26 @@ async def test_pre_call_forces_non_stream_for_loop():
)
@pytest.mark.asyncio
@pytest.mark.parametrize("call_type", (CallTypes.acompletion, CallTypes.aresponses))
async def test_pre_call_moves_stream_options_off_the_forced_non_stream_call(call_type: CallTypes):
logger: Final = CodeInterpreterInterceptionLogger(sandbox_config=FakeSandbox())
stream_options: Final = {"include_usage": True}
kwargs: Final[dict[str, object]] = {
"tools": [{"type": "code_interpreter", "container": {"type": "auto"}}],
"custom_llm_provider": "azure",
"stream": True,
"stream_options": stream_options,
}
out: Final = await logger.async_pre_call_deployment_hook(kwargs, call_type)
assert out is not None
assert out["stream"] is False
assert "stream_options" not in out
assert out[CODE_INTERPRETER_STREAM_OPTIONS_KEY] == stream_options
async def _build_plan(logger, sandbox, call_id="k1", provider="openai"):
return await logger.async_build_agentic_loop_plan(
tools={

14
tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py
View file

@ -1204,8 +1204,10 @@ def test_speech_response_without_a_byte_count_produces_no_output() -> None:
assert data.choices_out == ()
def test_speech_binary_response_is_logged_as_its_summary_not_dropped() -> None:
import httpx
@pytest.mark.parametrize("http_module", ("httpx", "httpx2"))
def test_speech_binary_response_is_logged_as_its_summary_not_dropped(http_module: str) -> None:
httpx: Final = pytest.importorskip(http_module)
from openai import HttpxBinaryResponseContent as SDKBinaryResponse
from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params
from litellm.litellm_core_utils.litellm_logging import _extract_response_obj_and_hidden_params
@ -1216,13 +1218,17 @@ def test_speech_binary_response_is_logged_as_its_summary_not_dropped() -> None:
set_provider_response_headers_in_hidden_params(speech, raw.headers)
response_obj, hidden_params = _extract_response_obj_and_hidden_params(speech, None)
assert isinstance(speech, SDKBinaryResponse)
assert speech.response is raw
assert speech.read() == raw.content
assert response_obj == {"object": "binary", "content_type": "audio/mpeg", "num_bytes": 1234}
assert hidden_params is not None
assert hidden_params["headers"]["content-type"] == "audio/mpeg"
def test_speech_binary_response_still_streaming_reports_the_bytes_downloaded_so_far() -> None:
import httpx
@pytest.mark.parametrize("http_module", ("httpx", "httpx2"))
def test_speech_binary_response_still_streaming_reports_the_bytes_downloaded_so_far(http_module: str) -> None:
httpx: Final = pytest.importorskip(http_module)
from litellm.types.llms.openai import HttpxBinaryResponseContent

69
tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py
View file

@ -4,6 +4,8 @@ Unit tests for WebSearch Interception Handler
Tests the WebSearchInterceptionLogger class and helper functions.
"""
from collections.abc import Awaitable, Callable
from typing import Any, Final
from unittest.mock import AsyncMock, MagicMock, Mock
import pytest
@ -20,7 +22,7 @@ from litellm.proxy._types import (
ProxyException,
UserAPIKeyAuth,
)
from litellm.types.utils import LlmProviders
from litellm.types.utils import CallTypes, LlmProviders
def test_initialize_from_proxy_config():
@ -810,6 +812,71 @@ async def test_deployment_hook_converts_stream_and_logging_obj_syncs():
assert logging_obj.stream is False
_STREAM_USAGE_OPTIONS = {"include_usage": True}
_Convert = Callable[[WebSearchInterceptionLogger, dict[str, Any]], Awaitable[dict[str, Any] | None]]
async def _convert_via_deployment_hook_chat(
logger: WebSearchInterceptionLogger, kwargs: dict[str, Any]
) -> dict[str, Any] | None:
return await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion)
async def _convert_via_deployment_hook_responses(
logger: WebSearchInterceptionLogger, kwargs: dict[str, Any]
) -> dict[str, Any] | None:
return await logger.async_pre_call_deployment_hook(
kwargs={**kwargs, "tools": [{"type": "web_search"}]}, call_type=CallTypes.aresponses
)
async def _convert_via_pre_request_hook(
logger: WebSearchInterceptionLogger, kwargs: dict[str, Any]
) -> dict[str, Any] | None:
return await logger.async_pre_request_hook(model=kwargs["model"], messages=kwargs["messages"], kwargs=kwargs)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"convert",
(_convert_via_deployment_hook_chat, _convert_via_deployment_hook_responses, _convert_via_pre_request_hook),
ids=("deployment_hook_chat", "deployment_hook_responses", "pre_request_hook"),
)
async def test_converted_stream_moves_stream_options_out_of_the_provider_request(convert: _Convert):
"""Every site that flips a streamed web search request to stream=False must take
stream_options with it: a provider that validates its body (Azure chat completions)
rejects stream_options on a non-streaming request, so the converted kwargs carry no
stream_options and stash the client's value for the replayed stream instead."""
logger: Final = WebSearchInterceptionLogger(enabled_providers=["azure"])
kwargs: Final = {
"model": "azure/gpt-4o",
"messages": [{"role": "user", "content": "Search for LiteLLM"}],
"custom_llm_provider": "azure",
"litellm_params": {"custom_llm_provider": "azure"},
"tools": [
{
"type": "function",
"function": {"name": LITELLM_WEB_SEARCH_TOOL_NAME, "parameters": {"type": "object"}},
}
],
"stream": True,
"stream_options": dict(_STREAM_USAGE_OPTIONS),
}
result: Final = await convert(logger, kwargs)
assert result is not None
assert (
result["stream"],
"stream_options" in result,
result["_websearch_interception_converted_stream"],
result["_websearch_interception_stream_options"],
) == (False, False, True, _STREAM_USAGE_OPTIONS)
def test_sync_forced_tool_choice_repoints_converted_web_search():
"""Regression (tool_choice 400): a forced tool_choice naming the original
web_search tool must be repointed to litellm_web_search after conversion.

78
tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py
View file

@ -20,7 +20,8 @@ removed, so `test_internal_control_fields_never_leak_into_provider_body` proves
they stay out of the body even without it.
"""
from typing import Any, Dict, List, Optional, Tuple
import datetime
from typing import Any, Dict, Final, List, Optional, Tuple
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -34,7 +35,12 @@ from litellm.integrations.code_interpreter_interception.handler import (
from litellm.litellm_core_utils.chat_completion_agentic_loop import (
maybe_run_chat_completion_agentic_loop,
)
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.integrations.custom_logger import (
CODE_INTERPRETER_INTERCEPTION_PREFIX,
HEADROOM_INTERCEPTION_PREFIX,
WEBSEARCH_INTERCEPTION_PREFIX,
AgenticLoopPlan,
AgenticLoopRequestPatch,
)
@ -253,6 +259,21 @@ class _GateOnlyLogger(CustomLogger):
self.cleanup_calls += 1
class _RecordingGateLogger(_GateOnlyLogger):
"""Also keeps every follow-up response the loop hands to the post hook."""
def __init__(self, plan: AgenticLoopPlan, tool_calls: Dict[str, Any]) -> None:
super().__init__(plan=plan, tool_calls=tool_calls)
self.followups: tuple[ModelResponse, ...] = ()
async def async_post_agentic_loop_response_hook(
self, response: object, plan: AgenticLoopPlan, kwargs: Dict[str, Any]
) -> object:
assert isinstance(response, ModelResponse)
self.followups = (*self.followups, response)
return response
def _patched_messages() -> List[Dict[str, Any]]:
return [
{"role": "user", "content": "what is 6*7?"},
@ -443,3 +464,58 @@ async def test_dispatcher_raises_on_repeated_tool_call_fingerprint(restore_callb
)
acompletion_mock.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("include_usage", (True, False), ids=("usage_requested", "usage_not_requested"))
@pytest.mark.parametrize(
"prefix", (WEBSEARCH_INTERCEPTION_PREFIX, CODE_INTERPRETER_INTERCEPTION_PREFIX, HEADROOM_INTERCEPTION_PREFIX)
)
async def test_converted_stream_replays_the_final_answer_with_the_requested_usage_chunk(
include_usage: bool,
prefix: str,
):
"""A streamed web search request runs non-streaming and is replayed as a fake
stream at depth 0. The stream_options the interception moved out of the provider
request decide whether the replay ends with the usage-only chunk the client asked for."""
plan: Final = AgenticLoopPlan(
run_agentic_loop=True,
request_patch=AgenticLoopRequestPatch(messages=_patched_messages()),
)
gate: Final = _RecordingGateLogger(plan=plan, tool_calls={"tool_calls": [{"id": "call_abc"}]})
logging_obj: Final = Logging(
model="gpt-4o",
messages=[{"role": "user", "content": "what is 6*7?"}],
stream=False,
call_type="acompletion",
start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc),
litellm_call_id="call-websearch",
function_id="fn-websearch",
dynamic_success_callbacks=[gate],
)
stash: Final = {f"{prefix}_stream_options": {"include_usage": True}} if include_usage else {}
result: Final = await maybe_run_chat_completion_agentic_loop(
response=_tool_call_model_response(),
model="gpt-4o",
messages=[{"role": "user", "content": "what is 6*7?"}],
optional_params={},
kwargs={f"{prefix}_converted_stream": True, "mock_response": "done", **stash},
logging_obj=logging_obj,
custom_llm_provider="azure",
stream=False,
)
assert isinstance(result, CustomStreamWrapper)
chunks: Final = [chunk async for chunk in result]
content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
usages: Final = [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None]
(followup,) = gate.followups
assert content == "done"
assert followup.usage.total_tokens > 0
assert [(usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) for usage in usages] == (
[(followup.usage.prompt_tokens, followup.usage.completion_tokens, followup.usage.total_tokens)]
if include_usage
else []
)

101
tests/unit/litellm_core_utils/test_realtime_streaming.py
View file

@ -1,6 +1,6 @@
import asyncio
import json
from collections.abc import Coroutine, Mapping
from collections.abc import Coroutine, Mapping, Sequence
from dataclasses import dataclass
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
@ -17,8 +17,11 @@ from litellm.litellm_core_utils.realtime_streaming import (
RealTimeStreaming,
client_sent_openai_beta_realtime_header,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.llms.xai.realtime.transformation import XAIRealtimeNormalizer
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.realtime import RealtimeResponseTransformInput, RealtimeResponseTypedDict
from litellm.types.utils import GenericGuardrailAPIInputs
@ -903,37 +906,75 @@ async def test_transcription_session_captures_usage_and_skips_response_create():
)
@pytest.mark.asyncio
async def test_non_transcription_completed_event_still_triggers_response_create():
"""
Regression guard: a normal (non-transcription) session with no guardrails must
keep triggering response.create on a completed transcription event.
"""
client_ws = MagicMock()
client_ws.send_text = AsyncMock()
class _PassthroughRealtimeConfig(BaseRealtimeConfig):
def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict:
return headers
completed = json.dumps(
{
"type": "conversation.item.input_audio_transcription.completed",
"transcript": "hi",
"item_id": "item_1",
def get_complete_url(self, api_base: str | None, model: str, api_key: str | None = None) -> str:
return api_base or ""
def transform_realtime_request(
self, message: str, model: str, session_configuration_request: str | None = None
) -> Sequence[str | bytes]:
return (message,)
def transform_realtime_response(
self,
message: str | bytes,
model: str,
logging_obj: LiteLLMLoggingObj,
realtime_response_transform_input: RealtimeResponseTransformInput,
) -> RealtimeResponseTypedDict:
return {
"response": [json.loads(message)],
"current_output_item_id": realtime_response_transform_input["current_output_item_id"],
"current_response_id": realtime_response_transform_input["current_response_id"],
"current_delta_chunks": realtime_response_transform_input["current_delta_chunks"],
"current_conversation_id": realtime_response_transform_input["current_conversation_id"],
"current_item_chunks": realtime_response_transform_input["current_item_chunks"],
"current_delta_type": realtime_response_transform_input["current_delta_type"],
"session_configuration_request": realtime_response_transform_input["session_configuration_request"],
}
).encode()
backend_ws = MagicMock()
backend_ws.recv = AsyncMock(side_effect=[completed, ConnectionClosed(None, None)])
backend_ws.send = AsyncMock()
logging_obj = MagicMock()
logging_obj.async_success_handler = AsyncMock()
logging_obj.success_handler = MagicMock()
@pytest.mark.asyncio
@pytest.mark.parametrize("through_provider_config", [False, True], ids=["raw", "provider_config"])
@pytest.mark.parametrize(
"event_hook,expected_response_creates",
[
(None, 0),
(GuardrailEventHooks.pre_call, 0),
(GuardrailEventHooks.post_call, 0),
(GuardrailEventHooks.realtime_input_transcription, 1),
],
ids=["no_guardrail", "pre_call", "post_call", "realtime_input_transcription"],
)
async def test_transcript_triggers_response_create_only_when_a_transcript_guardrail_owns_the_turn(
monkeypatch: pytest.MonkeyPatch,
through_provider_config: bool,
event_hook: GuardrailEventHooks | None,
expected_response_creates: int,
) -> None:
monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail(event_hook)] if event_hook else [])
completed: Final = _make_transcript_event("what are the opening hours tomorrow")
backend_ws: Final = MagicMock(
recv=AsyncMock(side_effect=[completed, ConnectionClosed(None, None)]),
send=AsyncMock(),
)
client_ws: Final = MagicMock(send_text=AsyncMock())
streaming: Final = RealTimeStreaming(
client_ws,
backend_ws,
MagicMock(async_success_handler=AsyncMock()),
provider_config=_PassthroughRealtimeConfig() if through_provider_config else None,
)
streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj)
await streaming.backend_to_client_send_messages()
assert streaming._is_transcription_session is False
sent_to_backend = [json.loads(c.args[0]) for c in backend_ws.send.call_args_list if c.args]
assert any(e.get("type") == "response.create" for e in sent_to_backend)
assert [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] == [
{"type": "response.create"}
] * expected_response_creates
assert [json.loads(call.args[0]) for call in client_ws.send_text.await_args_list] == [json.loads(completed)]
def test_client_session_update_does_not_mark_transcription_session():
@ -2865,10 +2906,10 @@ async def test_deferred_setup_clear_drops_appends_when_buffered():
assert streaming._pending_messages_until_setup == [new_audio]
def _transcription_guardrail():
"""A minimal real CustomGuardrail registered for the realtime transcript hook."""
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.types.guardrails import GuardrailEventHooks
def _transcription_guardrail(
event_hook: GuardrailEventHooks = GuardrailEventHooks.realtime_input_transcription,
) -> CustomGuardrail:
"""A minimal real CustomGuardrail registered for one realtime event hook."""
class _TranscriptionGuardrail(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
@ -2876,7 +2917,7 @@ def _transcription_guardrail():
return _TranscriptionGuardrail(
guardrail_name="test_transcription_guard",
event_hook=GuardrailEventHooks.realtime_input_transcription,
event_hook=event_hook,
default_on=True,
)

17
tests/unit/llms/anthropic/pass_through/adapters/test_handler_output_config_passthrough.py
View file

@ -27,6 +27,7 @@ Tests cover (consolidating PRs #23706 and #22727):
import os
import sys
from typing import Final
from unittest.mock import MagicMock, patch
import pytest
@ -40,6 +41,7 @@ from litellm.llms.anthropic.pass_through.adapters.handler import (
ANTHROPIC_ONLY_REQUEST_KEYS,
LiteLLMMessagesToCompletionTransformationHandler,
)
from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS
MESSAGES = [{"role": "user", "content": "hello"}]
@ -246,3 +248,18 @@ class TestPromptCacheOptionsForwarded:
result = _call_prepare(extra_kwargs={"prompt_cache_options": {"mode": "explicit"}}, model="gpt-5.6")
completion_kwargs = result[0] if isinstance(result, tuple) else result
assert completion_kwargs["prompt_cache_options"] == {"mode": "explicit"}
@pytest.mark.asyncio
@pytest.mark.parametrize("converted_stream_key", sorted(CONVERTED_STREAM_KEYS))
async def test_messages_layer_converted_stream_flag_still_gets_a_message_back(converted_stream_key: str):
response: Final = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler(
max_tokens=16,
messages=MESSAGES,
model="openai/gpt-4o",
stream=False,
mock_response="hi",
**{converted_stream_key: True},
)
assert response["content"][0]["text"] == "hi"

14
tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py
View file

@ -2,6 +2,7 @@ import datetime
import json
import os
import sys
from typing import Final
from unittest.mock import AsyncMock, patch
import pytest
@ -14,6 +15,7 @@ from litellm.llms.anthropic.pass_through.responses_adapters.handler import (
LiteLLMMessagesToResponsesAPIHandler,
_build_responses_kwargs,
)
from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS, converted_stream_requested
MESSAGES = [{"role": "user", "content": "hello"}]
CLAUDE_CODE_USER_ID = json.dumps({"device_id": "d" * 64, "account_uuid": "", "session_id": "session-abc"})
@ -31,6 +33,18 @@ RESPONSES_SSE_BODY = (
)
def test_build_responses_kwargs_never_asks_the_nested_call_to_replay_a_stream():
responses_kwargs: Final = _build_responses_kwargs(
max_tokens=1024,
messages=MESSAGES,
model="azure/gpt-5.6-luna",
extra_kwargs={"custom_llm_provider": "azure", **{key: True for key in CONVERTED_STREAM_KEYS}},
)
assert not converted_stream_requested(responses_kwargs)
assert responses_kwargs["custom_llm_provider"] == "azure"
def test_build_responses_kwargs_derives_prompt_cache_key_from_claude_code_session_id():
responses_kwargs = _build_responses_kwargs(
max_tokens=1024,

88
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -5,6 +5,8 @@ import json
import logging
import threading
import time
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from unittest.mock import AsyncMock, Mock, patch
@ -18,6 +20,9 @@ from litellm.integrations.code_interpreter_interception.handler import (
CodeInterpreterInterceptionLogger,
LITELLM_CODE_EXECUTION_TOOL_NAME,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger
from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
from litellm.llms.base_llm.audio_transcription.transformation import (
AudioTranscriptionRequestData,
BaseAudioTranscriptionConfig,
@ -353,6 +358,89 @@ def test_response_api_handler_runs_responses_pre_call_hook_before_transform():
assert hook_litellm_params.get(_SANDBOX_KEY)
class _AgenticLoopKwargsRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.seen_kwargs: Mapping[str, object] = MappingProxyType({})
async def async_should_run_agentic_loop(
self,
response: object,
model: str,
messages: list[dict[str, object]],
tools: list[dict[str, object]] | None,
stream: bool,
custom_llm_provider: str,
kwargs: dict[str, object],
) -> tuple[bool, dict[str, object]]:
self.seen_kwargs = MappingProxyType(dict(kwargs))
return False, {}
class _SentRequestRecorder:
def __init__(self) -> None:
self.body: Mapping[str, object] = MappingProxyType({})
def respond(self, request: httpx.Request) -> httpx.Response:
self.body = MappingProxyType(json.loads(request.content))
return httpx.Response(
200,
json={
"id": "resp_1",
"object": "response",
"created_at": 0,
"status": "completed",
"model": "gpt-4o",
"output": [
{
"type": "message",
"id": "msg_1",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "hello", "annotations": []}],
}
],
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
},
)
def test_response_api_handler_sends_no_stream_options_when_web_search_forces_a_non_streaming_call(
monkeypatch: pytest.MonkeyPatch,
):
sent: Final = _SentRequestRecorder()
recorder: Final = _AgenticLoopKwargsRecorder()
monkeypatch.setattr(
litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["azure"]), recorder]
)
logging_obj: Final = Mock()
logging_obj.dynamic_success_callbacks = []
BaseLLMHTTPHandler().response_api_handler(
model="gpt-4o",
input="what is new in litellm?",
responses_api_provider_config=AzureOpenAIResponsesAPIConfig(),
response_api_optional_request_params={
"stream": True,
"stream_options": {"include_obfuscation": True},
"tools": [{"type": "web_search"}],
},
custom_llm_provider="azure",
litellm_params=GenericLiteLLMParams(
api_key="sk-test",
api_base="https://example.openai.azure.com",
api_version="2025-04-01-preview",
stream_options={"include_obfuscation": True},
),
logging_obj=logging_obj,
client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(sent.respond))),
)
assert sent.body["stream"] is False
assert "stream_options" not in sent.body
assert "stream_options" not in recorder.seen_kwargs
@pytest.mark.asyncio
async def test_async_response_api_handler_streams_when_provider_transform_adds_stream():
handler = BaseLLMHTTPHandler()

130
tests/unit/llms/openai/test_openai.py
View file

@ -1,11 +1,16 @@
import asyncio, importlib, os
import asyncio
import importlib
import os
import json
from typing import Final
from unittest.mock import AsyncMock, Mock, patch
from collections.abc import Mapping
from types import ModuleType
from typing import Final, Protocol
from unittest.mock import Mock, patch
import httpx
import pytest
from openai import AsyncOpenAI, OpenAI
from openai._types import Response as SDKResponse
import litellm
from litellm.llms.openai.openai import (
@ -77,19 +82,48 @@ def test_get_stream_options_passes_caller_stream_options_through_on_any_host(api
}
@pytest.fixture(params=("httpx", "sdk"))
def http_backend(request: pytest.FixtureRequest) -> ModuleType:
return importlib.import_module("httpx" if request.param == "httpx" else SDKResponse.__module__)
@pytest.fixture(params=("openai", "perplexity", "cerebras", "nvidia_nim"))
def sdk_provider(request: pytest.FixtureRequest) -> str:
return str(request.param)
class _Request(Protocol):
@property
def content(self) -> bytes: ...
@property
def extensions(self) -> Mapping[str, object]: ...
@property
def url(self) -> object: ...
@property
def headers(self) -> Mapping[str, str]: ...
@pytest.mark.asyncio
async def test_acompletion_returns_json_reply_over_injected_transport():
async def test_acompletion_returns_json_reply_over_injected_transport(
http_backend: ModuleType, sdk_provider: str, monkeypatch: pytest.MonkeyPatch
):
outbound: Final = asyncio.Queue()
def respond(request: httpx.Request) -> httpx.Response:
def respond(request: _Request) -> httpx.Response | SDKResponse:
assert str(request.url) == f"https://{sdk_provider}.example/v1/chat/completions"
assert request.headers["authorization"] == f"Bearer {sdk_provider}-key"
assert request.extensions["timeout"] == {"connect": 7, "read": 7, "write": 7, "pool": 7}
outbound.put_nowait(json.loads(request.content))
return httpx.Response(
return http_backend.Response(
200,
json={
"id": "chatcmpl-smoke",
"object": "chat.completion",
"created": 0,
"model": "gpt-5.6",
"model": "sdk-compat",
"choices": [
{
"index": 0,
@ -101,31 +135,71 @@ async def test_acompletion_returns_json_reply_over_injected_transport():
},
)
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client:
client: Final = AsyncOpenAI(api_key="transport-only", http_client=http_client)
async with http_backend.AsyncClient(transport=http_backend.MockTransport(respond)) as http_client:
monkeypatch.setattr(litellm, "aclient_session", http_client)
response: Final = await asyncio.wait_for(
litellm.acompletion(
model="openai/gpt-5.6",
api_key="transport-only",
client=client,
model=f"{sdk_provider}/sdk-compat",
api_key=f"{sdk_provider}-key",
api_base=f"https://{sdk_provider}.example/v1",
messages=[{"role": "user", "content": "smoke-json-request"}],
timeout=7,
num_retries=0,
max_retries=0,
),
timeout=10,
)
request: Final = await asyncio.wait_for(outbound.get(), timeout=10)
assert request["model"] == "gpt-5.6"
assert request["model"] == "sdk-compat"
assert request["messages"] == [{"role": "user", "content": "smoke-json-request"}]
assert not request.get("stream")
assert outbound.empty()
assert response.choices[0].message.content == "smoke-json-reply"
assert response.choices[0].finish_reason == "stop"
assert response.usage.total_tokens == 15
assert not http_client.is_closed
@pytest.mark.asyncio
async def test_acompletion_streams_text_deltas_over_injected_transport():
@pytest.mark.parametrize(
("status", "expected_error"),
(
(400, litellm.BadRequestError),
(429, litellm.RateLimitError),
(503, litellm.ServiceUnavailableError),
(None, litellm.Timeout),
),
)
async def test_acompletion_preserves_public_errors_for_both_http_clients(
http_backend: ModuleType,
sdk_provider: str,
status: int | None,
expected_error: type[Exception],
monkeypatch: pytest.MonkeyPatch,
) -> None:
def respond(request: _Request) -> httpx.Response | SDKResponse:
if status is None:
raise http_backend.ReadTimeout("upstream timeout", request=request)
return http_backend.Response(status, json={"error": {"message": "upstream failure", "type": "api_error"}})
async with http_backend.AsyncClient(transport=http_backend.MockTransport(respond)) as http_client:
monkeypatch.setattr(litellm, "aclient_session", http_client)
with pytest.raises(expected_error, match="timed out" if status is None else "upstream failure"):
await litellm.acompletion(
model=f"{sdk_provider}/sdk-compat",
api_key=f"{sdk_provider}-key",
api_base=f"https://{sdk_provider}.example/v1",
messages=[{"role": "user", "content": "request"}],
num_retries=0,
max_retries=0,
)
assert not http_client.is_closed
@pytest.mark.asyncio
async def test_acompletion_streams_text_deltas_over_injected_transport(
http_backend: ModuleType, sdk_provider: str, monkeypatch: pytest.MonkeyPatch
):
outbound: Final = asyncio.Queue()
def chunk(delta: dict, finish: str | None) -> bytes:
@ -133,18 +207,18 @@ async def test_acompletion_streams_text_deltas_over_injected_transport():
"id": "chatcmpl-smoke",
"object": "chat.completion.chunk",
"created": 0,
"model": "gpt-5.6",
"model": "sdk-compat",
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
}
return f"data: {json.dumps(body)}\n\n".encode()
def respond(request: httpx.Request) -> httpx.Response:
def respond(request: _Request) -> httpx.Response | SDKResponse:
outbound.put_nowait(json.loads(request.content))
usage: Final = {
"id": "chatcmpl-smoke",
"object": "chat.completion.chunk",
"created": 0,
"model": "gpt-5.6",
"model": "sdk-compat",
"choices": [],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
}
@ -156,14 +230,14 @@ async def test_acompletion_streams_text_deltas_over_injected_transport():
b"data: [DONE]\n\n",
)
)
return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=content)
return http_backend.Response(200, headers={"content-type": "text/event-stream"}, content=content)
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client:
client: Final = AsyncOpenAI(api_key="transport-only", http_client=http_client)
async with http_backend.AsyncClient(transport=http_backend.MockTransport(respond)) as http_client:
monkeypatch.setattr(litellm, "aclient_session", http_client)
stream: Final = await litellm.acompletion(
model="openai/gpt-5.6",
api_key="transport-only",
client=client,
model=f"{sdk_provider}/sdk-compat",
api_key=f"{sdk_provider}-key",
api_base=f"https://{sdk_provider}.example/v1",
messages=[{"role": "user", "content": "smoke-stream-request"}],
stream=True,
num_retries=0,
@ -190,7 +264,7 @@ async def test_acompletion_streams_text_deltas_over_injected_transport():
@pytest.mark.asyncio
async def test_acompletion_streams_tool_call_arguments_over_injected_transport():
async def test_acompletion_streams_tool_call_arguments_over_injected_transport(http_backend: ModuleType):
outbound: Final = asyncio.Queue()
tools: Final = [
{
@ -217,7 +291,8 @@ async def test_acompletion_streams_tool_call_arguments_over_injected_transport()
}
return f"data: {json.dumps(body)}\n\n".encode()
def respond(request: httpx.Request) -> httpx.Response:
def respond(request: _Request) -> httpx.Response | SDKResponse:
assert request.extensions["timeout"] == {"connect": 1, "read": None, "write": 2, "pool": 3}
outbound.put_nowait(json.loads(request.content))
content: Final = b"".join(
(
@ -239,9 +314,9 @@ async def test_acompletion_streams_tool_call_arguments_over_injected_transport()
b"data: [DONE]\n\n",
)
)
return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=content)
return http_backend.Response(200, headers={"content-type": "text/event-stream"}, content=content)
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client:
async with http_backend.AsyncClient(transport=http_backend.MockTransport(respond)) as http_client:
client: Final = AsyncOpenAI(api_key="transport-only", http_client=http_client)
messages: Final = [{"role": "user", "content": "weather in Paris"}]
stream: Final = await litellm.acompletion(
@ -251,6 +326,7 @@ async def test_acompletion_streams_tool_call_arguments_over_injected_transport()
messages=messages,
tools=tools,
stream=True,
timeout=http_backend.Timeout(connect=1, read=None, write=2, pool=3),
num_retries=0,
max_retries=0,
)

Some files were not shown because too many files have changed in this diff Show more