diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 1cf79049a66..42b418b3ee0 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -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 diff --git a/ci_cd/cost_map_guard.py b/ci_cd/cost_map_guard.py index 5842cf6f1ac..14eecf01461 100644 --- a/ci_cd/cost_map_guard.py +++ b/ci_cd/cost_map_guard.py @@ -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: diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs index 1547a4daa93..5ad9f8df8b8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs @@ -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 { @@ -38,7 +38,7 @@ impl PythonBinding for ChatCompletionsPythonHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> Result> { - 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 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) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs index 2b69d0e6321..00ca5e092a8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs @@ -22,7 +22,7 @@ fn run_chat_completions( asynchronous: bool, ) -> PyResult> { let host = InferenceHost::new( - call.resolved()?.unbind(), + call.view()?, "litellm.rust_bridge.chat_completions.route_host", ); run_inference::( diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index 0333aaeb04a..bc442d1cf36 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -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, + view: RequestView, module: &'static str, } @@ -34,16 +39,17 @@ pub(super) struct ProjectedCall { } impl InferenceHost { - pub fn new(request: Py, 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 { + self.view.resolve(py, arguments)?; let argument = |name: &str| self.argument(py, arguments, name); let string = |name: &str| -> PyResult> { argument(name)?.map(|value| value.extract()).transpose() @@ -93,7 +99,15 @@ impl InferenceHost { arguments: &Bound<'py, PyDict>, name: &str, ) -> PyResult>> { - 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)) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index d34047c735d..adeba397b26 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -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 { /// 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, + view: RequestView, cache: PythonCache, } impl MessagesPythonHost { - pub(super) fn new(request: Py, 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> { - 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> { argument(name)?.map(|value| value.extract()).transpose() }; @@ -182,9 +182,9 @@ impl MessagesPythonHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> PyResult>> { - let request = self.request.bind(py); let mapping = |name: &str| -> PyResult>> { - 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> { - 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::>().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::>() @@ -264,6 +265,9 @@ impl PythonBinding for MessagesPythonHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> Result<(MessagesCall, Selection), InvokeError> { + 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> 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; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index 5f69cf078b7..c788760b268 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -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, ) diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 7427ec8da7d..416b94db55b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -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> { effective_py_args(&self.base, &self.kwargs) } + + /// The view a route host reads its arguments from. + fn view(&self) -> PyResult { + 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); + +struct RequestDicts { + base: Py, + request: Py, +} + +impl RequestView { + pub(crate) fn new(base: &Bound<'_, PyDict>, kwargs: &Bound<'_, PyDict>) -> PyResult { + 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>> { + 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> { @@ -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::() + .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::().unwrap()); + assert_eq!(value.as_deref(), expected); + }); + } + #[rstest] #[case::sync("transcription")] #[case::asynchronous("atranscription")] diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 5ae72acd104..dd3c648cc3e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -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, + view: RequestView, data: OcrHostData, } impl OcrPythonHost { - pub(super) fn new(request: Py) -> 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::>().map_err(PyErr::from)); match mapped { Ok(mapped) => PyErr::from_value(mapped.into_bound(py).into_any()), @@ -160,10 +161,11 @@ impl PythonHostCalls 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::() .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); diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index fab9ecc13c1..ec7f2ce79d9 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -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, ) diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs index b22bbc6192b..8261e66406b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs @@ -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 { @@ -59,7 +59,7 @@ impl PythonBinding for ResponsesPythonHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> Result> { - 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 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) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs index 50e5790a532..1bf24813103 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs @@ -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::(py, call, asynchronous, ResponsesPythonHost(host)) } diff --git a/litellm/__init__.py b/litellm/__init__.py index ac57a9d0fee..4196912488f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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", diff --git a/litellm/decisions/main.py b/litellm/decisions/main.py index 4a4af2b580e..f6abb1b2f3a 100644 --- a/litellm/decisions/main.py +++ b/litellm/decisions/main.py @@ -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 diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 439eadce13d..ac461a2472e 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -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 diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py index 764dd3ed3e0..69b528a29f8 100644 --- a/litellm/integrations/code_interpreter_interception/handler.py +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -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) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 15413140d9e..0fb482cbbd9 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -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 diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 9d7c9864e62..628486f377f 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -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, + ) diff --git a/litellm/litellm_core_utils/completion_timeout.py b/litellm/litellm_core_utils/completion_timeout.py index 163a4a6b9d6..e6f18257138 100644 --- a/litellm/litellm_core_utils/completion_timeout.py +++ b/litellm/litellm_core_utils/completion_timeout.py @@ -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) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 473be107cbb..c581db6370b 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -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 diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 731b39e681f..db43f7cd052 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -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(): diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py index 166ad916ea2..f20762789dd 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py @@ -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 ( diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index c640405c269..3564d7a1364 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -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, } diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index c57c643983a..724ca23c2ab 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -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, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a9b8fee2e8..a8462c55332 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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() diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index ffc5a5a71e0..839ca04e082 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -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 diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index 1bc752b6714..174aefd29e5 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -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, diff --git a/litellm/llms/openai/decisions/transformation.py b/litellm/llms/openai/decisions/transformation.py index 0afbe1a434e..317e5abeeb5 100644 --- a/litellm/llms/openai/decisions/transformation.py +++ b/litellm/llms/openai/decisions/transformation.py @@ -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" diff --git a/litellm/llms/openai/fine_tuning/handler.py b/litellm/llms/openai/fine_tuning/handler.py index c54263f4c35..851113bbd37 100644 --- a/litellm/llms/openai/fine_tuning/handler.py +++ b/litellm/llms/openai/fine_tuning/handler.py @@ -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 diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py index bc02d274f24..664d16c48bf 100644 --- a/litellm/llms/openai/image_variations/handler.py +++ b/litellm/llms/openai/image_variations/handler.py @@ -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, diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index a782bf35ac8..3a12326d433 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -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 diff --git a/litellm/llms/strands_decider/decisions/transformation.py b/litellm/llms/strands_decider/decisions/transformation.py index 9ef966e18f2..fd8a643be2a 100644 --- a/litellm/llms/strands_decider/decisions/transformation.py +++ b/litellm/llms/strands_decider/decisions/transformation.py @@ -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[a-z0-9-]+):[0-9]{12}:runtime/[A-Za-z0-9_-]+\Z" +) +_INVOCATIONS_PATH: Final = re.compile(r"\A/runtimes/(?P[^/]+)/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) diff --git a/litellm/main.py b/litellm/main.py index b80819354f0..4744a502194 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6d49e7cd7fd..400f153c115 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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": { diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 87dab4a8d67..1d86a20d80b 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -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]: """ diff --git a/litellm/proxy/_experimental/mcp_server/contracts.py b/litellm/proxy/_experimental/mcp_server/contracts.py index bf0fae1a210..c337c17c9ff 100644 --- a/litellm/proxy/_experimental/mcp_server/contracts.py +++ b/litellm/proxy/_experimental/mcp_server/contracts.py @@ -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)) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7c375969a44..e492af133ff 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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"] diff --git a/litellm/proxy/_experimental/mcp_server/operations.py b/litellm/proxy/_experimental/mcp_server/operations.py index d98748612bc..2f0b8e64c9e 100644 --- a/litellm/proxy/_experimental/mcp_server/operations.py +++ b/litellm/proxy/_experimental/mcp_server/operations.py @@ -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): diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 496c92f88ee..38f526bcd13 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -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: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index bf813eaff2b..17d3c343756 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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: diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 217ffcd6b5f..a8aff2b2e4b 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -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) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ff7f071d534..477edf47c5f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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=( diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 44f62fa2b26..6d820b0e801 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -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 diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4f18c773a24..7f0e0b1fc87 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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 diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py index b028e7fda20..089d0b7a226 100644 --- a/litellm/proxy/common_utils/proxy_rate_limit_error.py +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -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 diff --git a/litellm/proxy/decisions_endpoints/endpoints.py b/litellm/proxy/decisions_endpoints/endpoints.py index f2891c184bf..4f4209208ac 100644 --- a/litellm/proxy/decisions_endpoints/endpoints.py +++ b/litellm/proxy/decisions_endpoints/endpoints.py @@ -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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 36d48d49c7e..aa10a80f815 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index cfffff8eec1..421cf2281db 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -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}", diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 7c02ac308f3..d3ca4597f0d 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -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") diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 5b865f3b5c3..21025a8a820 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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 diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 490b393aa68..913310952cd 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 53f49f78c9f..dac46855fd6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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", diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index d2db9357349..13542f22f88 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -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. diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index ff31d7a65c4..096ad1c30dd 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -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) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index cd77d6088f6..26f65fee4fe 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -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: diff --git a/litellm/types/router.py b/litellm/types/router.py index 53df3ce6773..cbe332dab12 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -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 diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6d49e7cd7fd..400f153c115 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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": { diff --git a/packaging/litellm-core/pyproject.toml b/packaging/litellm-core/pyproject.toml index 9653e6f983a..cf87cc6b936 100644 --- a/packaging/litellm-core/pyproject.toml +++ b/packaging/litellm-core/pyproject.toml @@ -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", diff --git a/pyproject.toml b/pyproject.toml index 924e3a939b2..133fc735df1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/tests/base_sdk_tests/check_base_sdk_install.py b/tests/base_sdk_tests/check_base_sdk_install.py index e56d48f5e6e..3abea2a9707 100644 --- a/tests/base_sdk_tests/check_base_sdk_install.py +++ b/tests/base_sdk_tests/check_base_sdk_install.py @@ -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: diff --git a/tests/base_sdk_tests/check_sdk_http.py b/tests/base_sdk_tests/check_sdk_http.py index d773e252127..c7fb6a542ff 100644 --- a/tests/base_sdk_tests/check_sdk_http.py +++ b/tests/base_sdk_tests/check_sdk_http.py @@ -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" diff --git a/tests/integration/_support/forward_proxy.py b/tests/integration/_support/forward_proxy.py index 9f5cff94277..0d92c3d0719 100644 --- a/tests/integration/_support/forward_proxy.py +++ b/tests/integration/_support/forward_proxy.py @@ -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() diff --git a/tests/integration/_support/interception_vendor.py b/tests/integration/_support/interception_vendor.py new file mode 100644 index 00000000000..3ca40087441 --- /dev/null +++ b/tests/integration/_support/interception_vendor.py @@ -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 + ), + ) diff --git a/tests/integration/compatibility/test_openai_sdk_client_paths.py b/tests/integration/compatibility/test_openai_sdk_client_paths.py new file mode 100644 index 00000000000..f8d7449038a --- /dev/null +++ b/tests/integration/compatibility/test_openai_sdk_client_paths.py @@ -0,0 +1,1200 @@ +from __future__ import annotations + +import asyncio +import base64 +import binascii +import json +import os +import threading +import uuid +from collections.abc import Awaitable, Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, TypeVar + +import anthropic +import httpx +import openai +import pytest +import yaml +from anthropic.types import RawContentBlockDeltaEvent, RawMessageStartEvent, RawMessageStreamEvent, TextDelta +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.openai_wire import chat_reply, openai_error, responses_reply +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_vendor import response_identities +from integration._support.wire import Reply, Request, Wire, wire_server +from openai.types import Completion, ModerationCreateResponse +from openai.types.chat import ChatCompletionChunk +from openai.types.responses import ResponseStreamEvent +from pydantic import JsonValue + +from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + +PROXY_WORKERS: Final = int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1")) +WORKER_SYNC_SECONDS: Final = 0.0 if PROXY_WORKERS == 1 else PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + 5.0 +ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) +TEXT: Final = "openai sdk client path" +PROVIDER_KEY: Final = "integration-provider-key" +PROVIDER_AUTHORIZATION: Final = f"Bearer {PROVIDER_KEY}" +UPSTREAM_MODEL: Final = "gpt-4o-mini" +DEPLOYMENT_MODEL: Final = f"openai/{UPSTREAM_MODEL}" +INSTRUCT_MODEL: Final = "gpt-3.5-turbo-instruct" +INSTRUCT_DEPLOYMENT_MODEL: Final = f"openai/{INSTRUCT_MODEL}" +NO_CACHE: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache": {"no-cache": True}}) +USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8} +) +FIVE_KB: Final = "x" * 5120 +REJECTED_CHAT_TIMEOUTS: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"list": [1, 2], "dict": {"read": 5}, "5 KB string": FIVE_KB} +) +ACCEPTED_CHAT_TIMEOUTS: Final[Mapping[str, JsonValue]] = MappingProxyType({"empty string": "", "int": 30}) +NO_EXTRA: Final[Mapping[str, JsonValue]] = MappingProxyType({}) +REJECTED_DEPLOYMENT_TIMEOUTS: Final[Mapping[str, JsonValue]] = MappingProxyType({"list": [1, 2], "dict": {"read": 5}}) +ACCEPTED_DEPLOYMENT_TIMEOUTS: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"empty string": "", "5 KB string": FIVE_KB} +) +MODERATION_CATEGORIES: Final = ("hate", "harassment", "self-harm", "sexual", "violence") +LISTED_JOB: Final = "ftjob-listed" +ASSISTANT: Final = "asst_integration" +CLOUDFLARE_GATEWAY: Final = "/gateway.ai.cloudflare.com/v1/account/gateway/azure-openai/resource" +AZURE_API_VERSION: Final = "2024-10-21" +STREAM_HEAD: Final = b"h" * 65536 +STREAM_TAIL: Final = b'{"custom_id": "tail", "response": {"status_code": 200}}\n' +SPEECH_AUDIO: Final = b"OggS" + bytes(range(60)) +SPEECH_UPSTREAM_MODEL: Final = "gpt-4o-mini-tts" +PAYMENT_REQUIRED_DEPLOYMENTS: Final[Mapping[str, str]] = MappingProxyType( + {"openai": "openai/gpt-4o-mini", "anthropic": "anthropic/claude-sonnet-4-5"} +) +_R: Final = TypeVar("_R") + + +def _identity(prefix: str) -> str: + return f"{prefix}-{uuid.uuid4().hex}" + + +def _path(request: Request) -> str: + return request.target.split("?", 1)[0] + + +def _body(request: Request) -> Mapping[str, JsonValue]: + return JSON_OBJECT.validate_json(request.body) + + +def _streams(request: Request) -> bool: + return _body(request).get("stream") is True + + +def _upstream_model(request: Request) -> str: + return string_value(_body(request).get("model") or UPSTREAM_MODEL) + + +def _json( + payload: Mapping[str, JsonValue], *, status: int = 200, headers: Mapping[str, str] = MappingProxyType({}) +) -> Reply: + return Reply(status=status, body=json.dumps(payload).encode(), headers=headers) + + +def _peer(route: Callable[[Request], Reply | None]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET" and _path(request).endswith("/models"): + return _json({"object": "list", "data": []}) + reply: Final = route(request) + if reply is None: + return _json({"error": {"message": f"unrouted {request.method} {request.target}"}}, status=404) + return reply + + return respond + + +def _is_model_discovery(request: Request) -> bool: + return request.method == "GET" and _path(request).endswith("/models") + + +def _served(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if not _is_model_discovery(request)) + + +def _routes(served: tuple[Request, ...]) -> tuple[tuple[str, str], ...]: + return tuple((request.method, _path(request)) for request in served) + + +def _is_chat(request: Request) -> bool: + return request.method == "POST" and _path(request).endswith("/chat/completions") + + +def _chat_peer(identity: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if not _is_chat(request): + return None + return chat_reply(identity, _upstream_model(request), TEXT, stream=_streams(request)) + + return _peer(route) + + +def _fresh_chat_peer(served: SimpleQueue[str]) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if not _is_chat(request): + return None + identity: Final = _identity("chatcmpl") + served.put(identity) + return chat_reply(identity, _upstream_model(request), TEXT, stream=False) + + return _peer(route) + + +def _held_chat_peer(gate: threading.Event, identity: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if not _is_chat(request): + return None + assert gate.wait(timeout=10), "the timeout cell never released its peer" + return chat_reply(identity, _upstream_model(request), TEXT, stream=False) + + return _peer(route) + + +def _status_peer(reply: Reply) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + return reply if request.method == "POST" else None + + return _peer(route) + + +def _completion_frame(identity: str, model: str, text: str, finish_reason: str | None) -> bytes: + frame: Final = { + "id": identity, + "object": "text_completion", + "created": 1, + "model": model, + "choices": [{"text": text, "index": 0, "finish_reason": finish_reason, "logprobs": None}], + } + return b"data: " + json.dumps(frame).encode() + b"\n\n" + + +def _completion_reply(identity: str, model: str, *, stream: bool) -> Reply: + if stream: + return Reply( + content_type="text/event-stream", + chunks=( + _completion_frame(identity, model, TEXT, None), + _completion_frame(identity, model, "", "stop") + b"data: [DONE]\n\n", + ), + ) + return _json( + { + "id": identity, + "object": "text_completion", + "created": 1, + "model": model, + "choices": [{"text": TEXT, "index": 0, "finish_reason": "stop", "logprobs": None}], + "usage": dict(USAGE), + } + ) + + +def _completion_peer(identity: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if request.method != "POST" or _path(request) != "/v1/completions": + return None + return _completion_reply(identity, _upstream_model(request), stream=_streams(request)) + + return _peer(route) + + +def _responses_peer(identity: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if request.method != "POST" or _path(request) != "/v1/responses": + return None + return responses_reply(identity, _upstream_model(request), TEXT, stream=_streams(request)) + + return _peer(route) + + +def _moderation(identity: str, model: str) -> Mapping[str, JsonValue]: + return { + "id": identity, + "model": model, + "results": [ + { + "flagged": False, + "categories": {category: False for category in MODERATION_CATEGORIES}, + "category_scores": {category: 0.0 for category in MODERATION_CATEGORIES}, + "category_applied_input_types": {category: ["text"] for category in MODERATION_CATEGORIES}, + } + ], + } + + +def _moderation_peer(identity: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if request.method != "POST" or _path(request) != "/v1/moderations": + return None + return _json(_moderation(identity, _upstream_model(request))) + + return _peer(route) + + +def _batch_line(custom_id: str) -> bytes: + line: Final = { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": UPSTREAM_MODEL, "messages": [_user_message()]}, + } + return json.dumps(line).encode() + b"\n" + + +def _file_object(file_id: str, size: int) -> Mapping[str, JsonValue]: + return { + "id": file_id, + "object": "file", + "purpose": "batch", + "bytes": size, + "created_at": 1, + "filename": "batch.jsonl", + "status": "processed", + } + + +def _files_peer(file_id: str, content: bytes, chunks: tuple[bytes, ...] | None = None) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + path: Final = _path(request) + if request.method == "POST" and path == "/v1/files": + return _json(_file_object(file_id, len(content))) + if request.method == "GET" and path == f"/v1/files/{file_id}": + return _json(_file_object(file_id, len(content))) + if request.method == "GET" and path == f"/v1/files/{file_id}/content": + return Reply(content_type="application/octet-stream", body=content, chunks=chunks) + return None + + return _peer(route) + + +def _batch_object(batch_id: str, input_file_id: str, status: str) -> Mapping[str, JsonValue]: + return { + "id": batch_id, + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": input_file_id, + "completion_window": "24h", + "status": status, + "output_file_id": None, + "error_file_id": None, + "created_at": 1, + "in_progress_at": None, + "expires_at": None, + "finalizing_at": None, + "completed_at": None, + "failed_at": None, + "expired_at": None, + "cancelling_at": None, + "cancelled_at": None, + "request_counts": {"total": 1, "completed": 0, "failed": 0}, + "metadata": None, + } + + +def _batches_peer(file_id: str, batch_id: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + path: Final = _path(request) + if request.method == "POST" and path == "/v1/files": + return _json(_file_object(file_id, 1)) + if request.method == "POST" and path == "/v1/batches": + return _json(_batch_object(batch_id, string_value(_body(request)["input_file_id"]), "validating")) + if request.method == "GET" and path == f"/v1/batches/{batch_id}": + return _json(_batch_object(batch_id, file_id, "in_progress")) + if request.method == "GET" and path == "/v1/batches": + return _json( + { + "object": "list", + "data": [_batch_object(batch_id, file_id, "in_progress")], + "first_id": batch_id, + "last_id": batch_id, + "has_more": False, + } + ) + if request.method == "POST" and path == f"/v1/batches/{batch_id}/cancel": + return _json(_batch_object(batch_id, file_id, "cancelling")) + return None + + return _peer(route) + + +def _fine_tuning_job(job_id: str, status: str) -> Mapping[str, JsonValue]: + return { + "id": job_id, + "object": "fine_tuning.job", + "model": UPSTREAM_MODEL, + "created_at": 1, + "fine_tuned_model": None, + "finished_at": None, + "hyperparameters": {"n_epochs": 1, "batch_size": 1, "learning_rate_multiplier": 1.0}, + "organization_id": "org-integration", + "result_files": [], + "status": status, + "trained_tokens": None, + "training_file": f"file-{job_id.removeprefix('ftjob-')}", + "validation_file": None, + "seed": 1, + "error": None, + } + + +def _assistant(identity: str) -> Mapping[str, JsonValue]: + return { + "id": identity, + "object": "assistant", + "created_at": 1, + "name": "integration", + "description": None, + "model": UPSTREAM_MODEL, + "instructions": None, + "tools": [], + "metadata": {}, + "temperature": 1.0, + "top_p": 1.0, + "response_format": "auto", + "tool_resources": None, + } + + +def _configured_peer() -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + path: Final = _path(request) + if request.method == "POST" and path == "/v1/fine_tuning/jobs": + training_file: Final = string_value(_body(request)["training_file"]) + return _json(_fine_tuning_job(f"ftjob-{training_file.removeprefix('file-')}", "queued")) + if request.method == "GET" and path == "/v1/fine_tuning/jobs": + return _json({"object": "list", "data": [_fine_tuning_job(LISTED_JOB, "running")], "has_more": False}) + if request.method == "POST" and path.startswith("/v1/fine_tuning/jobs/") and path.endswith("/cancel"): + return _json( + _fine_tuning_job(path.removeprefix("/v1/fine_tuning/jobs/").removesuffix("/cancel"), "cancelled") + ) + if request.method == "GET" and path.startswith("/v1/fine_tuning/jobs/"): + return _json(_fine_tuning_job(path.removeprefix("/v1/fine_tuning/jobs/"), "running")) + if request.method == "GET" and path == "/v1/assistants": + return _json( + { + "object": "list", + "data": [_assistant(ASSISTANT)], + "first_id": ASSISTANT, + "last_id": ASSISTANT, + "has_more": False, + } + ) + return None + + return _peer(route) + + +def _config_with_provider_settings(directory: Path, api_base: str) -> Path: + shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text())) + path: Final = directory / "proxy_config_with_provider_settings.yaml" + path.write_text( + yaml.safe_dump( + { + **shared, + "finetune_settings": [{"custom_llm_provider": "openai", "api_base": api_base, "api_key": PROVIDER_KEY}], + "assistant_settings": { + "custom_llm_provider": "openai", + "litellm_params": {"api_base": api_base, "api_key": PROVIDER_KEY}, + }, + } + ) + ) + return path + + +@dataclass(frozen=True, slots=True) +class ConfiguredProxy: + gateway: Gateway + wire: Wire + + +@pytest.fixture(scope="module") +def configured_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[ConfiguredProxy]: + with gateway_from_environment() as rig_gateway, wire_server(_configured_peer()) as wire: + directory: Final = tmp_path_factory.mktemp("openai_sdk_client_paths") + config: Final = _config_with_provider_settings(directory, f"{wire.url}/v1") + with owned_proxy(rig_gateway, directory, {}, config=config, workers=PROXY_WORKERS) as owned: + yield ConfiguredProxy(owned, wire) + + +def _v1(gateway: Gateway) -> str: + return f"{str(gateway.client.base_url).rstrip('/')}/v1" + + +def _origin(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +@contextmanager +def _openai(gateway: Gateway) -> Iterator[openai.OpenAI]: + with ( + httpx.Client(timeout=15, trust_env=False) as transport, + openai.OpenAI(api_key=gateway.key, base_url=_v1(gateway), max_retries=0, http_client=transport) as client, + ): + yield client + + +def _async_openai(gateway: Gateway, call: Callable[[openai.AsyncOpenAI], Awaitable[_R]]) -> _R: + async def run() -> _R: + async with ( + httpx.AsyncClient(timeout=15, trust_env=False) as transport, + openai.AsyncOpenAI( + api_key=gateway.key, base_url=_v1(gateway), max_retries=0, http_client=transport + ) as client, + ): + return await call(client) + + return asyncio.run(run()) + + +@contextmanager +def _anthropic(gateway: Gateway) -> Iterator[anthropic.Anthropic]: + with ( + httpx.Client(timeout=15, trust_env=False) as transport, + anthropic.Anthropic( + api_key=gateway.key, base_url=_origin(gateway), max_retries=0, http_client=transport + ) as client, + ): + yield client + + +def _async_anthropic(gateway: Gateway, call: Callable[[anthropic.AsyncAnthropic], Awaitable[_R]]) -> _R: + async def run() -> _R: + async with ( + httpx.AsyncClient(timeout=15, trust_env=False) as transport, + anthropic.AsyncAnthropic( + api_key=gateway.key, base_url=_origin(gateway), max_retries=0, http_client=transport + ) as client, + ): + return await call(client) + + return asyncio.run(run()) + + +def _user_message() -> Mapping[str, JsonValue]: + return {"role": "user", "content": TEXT} + + +def _chat_body(model: str, extra: Mapping[str, JsonValue] = NO_EXTRA) -> Mapping[str, JsonValue]: + return {"model": model, "messages": [_user_message()], **NO_CACHE, **extra} + + +def _duplicated_timeout_body(model: str) -> bytes: + fields: Final = json.dumps(_chat_body(model)).removesuffix("}") + return f'{fields}, "timeout": 30, "timeout": 30}}'.encode() + + +def _chat(gateway: Gateway, model: str, extra: Mapping[str, JsonValue] = NO_EXTRA) -> httpx.Response: + return gateway.request("POST", "/v1/chat/completions", _chat_body(model, extra)) + + +def _response_id(response: httpx.Response) -> str: + assert response.status_code == 200, response.text + return string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + +def _error_message(response: httpx.Response) -> str: + body: Final = JSON_OBJECT.validate_json(response.content) + if "error" in body: + return string_value(object_value(body["error"])["message"]) + detail: Final = body["detail"] + if isinstance(detail, str): + return detail + assert isinstance(detail, list) and detail, response.text + return string_value(object_value(detail[0])["msg"]) + + +def _assert_rejected(response: httpx.Response, shape: str) -> None: + assert response.status_code >= 400, (shape, response.status_code, response.text) + assert _error_message(response), (shape, response.text) + + +def _assert_served(response: httpx.Response, identity: str, shape: str) -> None: + assert response.status_code == 200, (shape, response.status_code, response.text) + assert _response_id(response) == identity, (shape, response.text) + + +def _moderation_through_the_sdk(gateway: Gateway, model: str) -> ModerationCreateResponse | openai.APIStatusError: + with _openai(gateway) as client: + try: + return client.moderations.create(model=model, input=TEXT) + except openai.APIStatusError as error: + return error + + +def _moderation_after_worker_sync(gateway: Gateway, model: str) -> ModerationCreateResponse | openai.APIStatusError: + return eventually( + lambda: _moderation_through_the_sdk(gateway, model), + lambda outcome: not isinstance(outcome, openai.InternalServerError), + seconds=WORKER_SYNC_SECONDS + 10, + ) + + +def _speech_peer(content_type: str) -> Callable[[Request], Reply]: + def route(request: Request) -> Reply | None: + if request.method != "POST" or not _path(request).endswith("/audio/speech"): + return None + return Reply(body=SPEECH_AUDIO, content_type=content_type) + + return _peer(route) + + +def _speech_through_the_sdk(gateway: Gateway, model: str) -> httpx.Response | openai.APIStatusError: + with _openai(gateway) as client: + try: + return client.audio.speech.create(model=model, voice="alloy", input=TEXT, response_format="mp3").response + except openai.APIStatusError as error: + return error + + +def _speech_after_worker_sync(gateway: Gateway, model: str) -> httpx.Response | openai.APIStatusError: + return eventually( + lambda: _speech_through_the_sdk(gateway, model), + lambda outcome: not isinstance(outcome, openai.InternalServerError), + seconds=WORKER_SYNC_SECONDS + 10, + ) + + +def _attempt(call: Callable[[], _R]) -> _R | openai.APIStatusError: + try: + return call() + except openai.APIStatusError as error: + return error + + +def _unknown_to_the_serving_worker(outcome: object) -> bool: + return isinstance(outcome, openai.BadRequestError) and "Invalid model name" in outcome.response.text + + +def _outcome_after_worker_sync(call: Callable[[], _R]) -> _R | openai.APIStatusError: + return eventually( + lambda: _attempt(call), + lambda outcome: not _unknown_to_the_serving_worker(outcome), + seconds=WORKER_SYNC_SECONDS + 10, + ) + + +def _after_worker_sync(call: Callable[[], _R]) -> _R: + outcome: Final = _outcome_after_worker_sync(call) + assert not isinstance(outcome, openai.APIStatusError), outcome.response.text + return outcome + + +def _async_after_worker_sync(gateway: Gateway, call: Callable[[openai.AsyncOpenAI], Awaitable[_R]]) -> _R: + return _after_worker_sync(lambda: _async_openai(gateway, call)) + + +def _streamed_file_content(gateway: Gateway, file_id: str, model: str) -> tuple[int, tuple[bytes, ...]]: + with gateway.client.stream( + "GET", + f"/v1/files/{file_id}/content", + params={"model": model}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + return response.status_code, tuple(response.iter_bytes()) + + +def _upstream_id(managed: str, prefix: str) -> str: + encoded: Final = managed.removeprefix(prefix) + try: + decoded: Final = base64.urlsafe_b64decode(encoded + "=" * (-len(encoded) % 4)).decode() + except (binascii.Error, UnicodeDecodeError): + return managed + return decoded.removeprefix("litellm:").split(";", 1)[0] if decoded.startswith("litellm:") else managed + + +def _spend_row(identity: str) -> Mapping[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, call_type, model FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (identity,) + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + return rows[0] + + +def _spend_row_keyed_by_any(identities: frozenset[str]) -> Mapping[str, JsonValue]: + keys: Final = tuple(sorted(identities)) + placeholders: Final = ", ".join("%s" for _ in keys) + rows: Final = eventually( + lambda: read_rows( + f'SELECT request_id, call_type, model FROM "LiteLLM_SpendLogs" WHERE request_id IN ({placeholders})', keys + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + return rows[0] + + +def _spend_request_ids(identities: tuple[str, str, str]) -> frozenset[str]: + rows: Final = eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id IN (%s, %s, %s)', identities), + lambda rows: len(rows) == 3, + seconds=70, + ) + return frozenset(string_value(row["request_id"]) for row in rows) + + +def _anthropic_text(event: RawMessageStreamEvent) -> str: + if isinstance(event, RawContentBlockDeltaEvent) and isinstance(event.delta, TextDelta): + return event.delta.text + return "" + + +def _chunk_text(chunk: ChatCompletionChunk) -> str: + return chunk.choices[0].delta.content or "" if chunk.choices else "" + + +def _chunk_finished(chunk: ChatCompletionChunk) -> bool: + return bool(chunk.choices) and chunk.choices[0].finish_reason == "stop" + + +def _new_model(gateway: Gateway, api_base: str, timeout: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/model/new", + { + "model_name": f"integration-{uuid.uuid4().hex}", + "litellm_params": { + "model": f"openai/{UPSTREAM_MODEL}", + "api_key": PROVIDER_KEY, + "api_base": api_base, + "timeout": timeout, + }, + "model_info": {}, + }, + ) + + +def _register_accepted_model(scenario: Scenario, response: httpx.Response, shape: str) -> None: + assert response.status_code == 200, (shape, response.status_code, response.text) + identity: Final = string_value(object_value(JSON_OBJECT.validate_json(response.content)["model_info"])["id"]) + scenario.cleanups.callback(scenario.delete_model, identity) + + +def test_r01_chat_completions_non_stream_through_the_sync_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("chatcmpl") + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + with _openai(gateway) as client: + response: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": TEXT}], extra_body=dict(NO_CACHE) + ) + assert response.id == identity, response + assert response.choices[0].message.content == TEXT, response + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/chat/completions"),), served + assert served[0].headers["authorization"] == PROVIDER_AUTHORIZATION, served[0].headers + upstream_body: Final = _body(served[0]) + assert upstream_body["model"] == UPSTREAM_MODEL, upstream_body + assert upstream_body["messages"] == [_user_message()], upstream_body + assert _spend_row(identity)["model"] == DEPLOYMENT_MODEL + + +def test_r02_chat_completions_stream_through_the_async_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("chatcmpl") + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + + async def stream(client: openai.AsyncOpenAI) -> tuple[ChatCompletionChunk, ...]: + chunks: Final = await client.chat.completions.create( + model=model, messages=[{"role": "user", "content": TEXT}], stream=True, extra_body=dict(NO_CACHE) + ) + return tuple([chunk async for chunk in chunks]) + + received: Final = _async_openai(gateway, stream) + assert {chunk.id for chunk in received} == {identity}, received + assert "".join(_chunk_text(chunk) for chunk in received) == TEXT, received + assert any(_chunk_finished(chunk) for chunk in received), received + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/chat/completions"),), served + assert _body(served[0])["stream"] is True, served[0].body + assert _spend_row(identity)["model"] == DEPLOYMENT_MODEL + + +def test_r03_chat_completions_through_raw_httpx_with_a_body_timeout(gateway: Gateway) -> None: + identity: Final = _identity("chatcmpl") + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + assert _response_id(_chat(gateway, model, {"timeout": 30})) == identity + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/chat/completions"),), served + assert "timeout" not in _body(served[0]), served[0].body + assert _spend_row(identity)["model"] == DEPLOYMENT_MODEL + + +def test_r04_messages_non_stream_through_the_sync_anthropic_sdk(gateway: Gateway) -> None: + identity: Final = _identity("resp") + with gateway.scenario() as scenario, wire_server(_responses_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + with _anthropic(gateway) as client: + message: Final = client.messages.create( + model=model, max_tokens=64, messages=[{"role": "user", "content": TEXT}] + ) + assert identity in response_identities(message.id), message.id + first_block: Final = message.content[0] + assert first_block.type == "text" and first_block.text == TEXT, message + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/responses"),), served + assert TEXT in served[0].body.decode(), served[0].body + assert _spend_row(message.id)["call_type"] == "anthropic_messages" + + +def test_r05_messages_stream_through_the_async_anthropic_sdk(gateway: Gateway) -> None: + identity: Final = _identity("resp") + with gateway.scenario() as scenario, wire_server(_responses_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + + async def stream(client: anthropic.AsyncAnthropic) -> tuple[RawMessageStreamEvent, ...]: + events: Final = await client.messages.create( + model=model, max_tokens=64, messages=[{"role": "user", "content": TEXT}], stream=True + ) + return tuple([event async for event in events]) + + received: Final = _async_anthropic(gateway, stream) + first: Final = received[0] + assert isinstance(first, RawMessageStartEvent), received + assert received[-1].type == "message_stop", received + assert "".join(_anthropic_text(event) for event in received) == TEXT, received + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/responses"),), served + assert _body(served[0])["stream"] is True, served[0].body + assert _spend_row(first.message.id)["call_type"] == "anthropic_messages" + + +def test_r06_responses_non_stream_through_the_sync_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("resp") + with gateway.scenario() as scenario, wire_server(_responses_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + with _openai(gateway) as client: + response: Final = client.responses.create(model=model, input=TEXT, extra_body=dict(NO_CACHE)) + assert identity in response_identities(response.id), response.id + assert response.output_text == TEXT, response + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/responses"),), served + upstream_body: Final = _body(served[0]) + assert upstream_body["model"] == UPSTREAM_MODEL, upstream_body + assert upstream_body["input"] == TEXT, upstream_body + assert _spend_row(response.id)["model"] == DEPLOYMENT_MODEL + + +def test_r07_responses_stream_through_the_async_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("resp") + with gateway.scenario() as scenario, wire_server(_responses_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + + async def stream(client: openai.AsyncOpenAI) -> tuple[ResponseStreamEvent, ...]: + events: Final = await client.responses.create( + model=model, input=TEXT, stream=True, extra_body=dict(NO_CACHE) + ) + return tuple([event async for event in events]) + + received: Final = _async_openai(gateway, stream) + assert received[0].type == "response.created", received + last: Final = received[-1] + assert last.type == "response.completed", received + assert identity in response_identities(last.response.id), last + assert last.response.output_text == TEXT, last + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/responses"),), served + assert _body(served[0])["stream"] is True, served[0].body + assert _spend_row_keyed_by_any(response_identities(last.response.id))["model"] == DEPLOYMENT_MODEL + + +def test_r08_completions_non_stream_through_the_sync_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("cmpl") + with gateway.scenario() as scenario, wire_server(_completion_peer(identity)) as wire: + model: Final = scenario.model(model=INSTRUCT_DEPLOYMENT_MODEL, api_base=f"{wire.url}/v1") + with _openai(gateway) as client: + completion: Final = client.completions.create(model=model, prompt=TEXT, extra_body=dict(NO_CACHE)) + assert completion.id == identity, completion + assert completion.choices[0].text == TEXT, completion + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/completions"),), served + upstream_body: Final = _body(served[0]) + assert upstream_body["model"] == INSTRUCT_MODEL, upstream_body + assert upstream_body["prompt"] == TEXT, upstream_body + assert _spend_row(identity)["model"] == INSTRUCT_DEPLOYMENT_MODEL + + +def test_r09_completions_stream_through_the_async_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("cmpl") + with gateway.scenario() as scenario, wire_server(_completion_peer(identity)) as wire: + model: Final = scenario.model(model=INSTRUCT_DEPLOYMENT_MODEL, api_base=f"{wire.url}/v1") + + async def stream(client: openai.AsyncOpenAI) -> tuple[Completion, ...]: + chunks: Final = await client.completions.create( + model=model, prompt=TEXT, stream=True, extra_body=dict(NO_CACHE) + ) + return tuple([chunk async for chunk in chunks]) + + received: Final = _async_openai(gateway, stream) + assert {chunk.id for chunk in received} == {identity}, received + assert "".join(chunk.choices[0].text or "" for chunk in received if chunk.choices) == TEXT, received + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/completions"),), served + assert _body(served[0])["stream"] is True, served[0].body + assert _spend_row(identity)["model"] == INSTRUCT_DEPLOYMENT_MODEL + + +def test_r10_moderations_through_the_sync_openai_sdk(gateway: Gateway) -> None: + identity: Final = _identity("modr") + with gateway.scenario() as scenario, wire_server(_moderation_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + moderation: Final = _moderation_after_worker_sync(gateway, model) + assert isinstance(moderation, ModerationCreateResponse), moderation + assert moderation.id == identity, moderation + assert [result.flagged for result in moderation.results] == [False], moderation + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/moderations"),), served + assert served[0].headers["authorization"] == PROVIDER_AUTHORIZATION, served[0].headers + assert _body(served[0])["input"] == TEXT, served[0].body + assert _spend_row(moderation.id)["call_type"] == "amoderation" + + +def test_r11_files_create_retrieve_and_content_through_the_sync_openai_sdk(gateway: Gateway) -> None: + file_id: Final = _identity("file") + content: Final = _batch_line(file_id) + with gateway.scenario() as scenario, wire_server(_files_peer(file_id, content)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + with _openai(gateway) as client: + created: Final = _after_worker_sync( + lambda: client.files.create(file=("batch.jsonl", content), purpose="batch", extra_body={"model": model}) + ) + retrieved: Final = _after_worker_sync( + lambda: client.files.retrieve(created.id, extra_query={"model": model}) + ) + downloaded: Final = _after_worker_sync( + lambda: client.files.content(created.id, extra_query={"model": model}) + ).content + assert _upstream_id(created.id, "file-") == file_id, created + assert _upstream_id(retrieved.id, "file-") == file_id, retrieved + assert retrieved.bytes == len(content), retrieved + assert downloaded == content, downloaded + served: Final = _served(wire) + assert _routes(served) == ( + ("POST", "/v1/files"), + ("GET", f"/v1/files/{file_id}"), + ("GET", f"/v1/files/{file_id}/content"), + ), served + assert content in served[0].body, served[0].body + assert {request.headers["authorization"] for request in served} == {PROVIDER_AUTHORIZATION}, served + + +def test_r12_file_content_streams_through_raw_httpx(gateway: Gateway) -> None: + file_id: Final = _identity("file") + content: Final = STREAM_HEAD + STREAM_TAIL + with ( + gateway.scenario() as scenario, + wire_server(_files_peer(file_id, content, chunks=(STREAM_HEAD, STREAM_TAIL))) as wire, + ): + model: Final = scenario.model(api_base=f"{wire.url}/v1") + status, received = eventually( + lambda: _streamed_file_content(gateway, file_id, model), + lambda outcome: not (outcome[0] == 400 and b"Invalid model name" in b"".join(outcome[1])), + seconds=WORKER_SYNC_SECONDS + 10, + ) + assert status == 200, b"".join(received).decode() + assert b"".join(received) == content, (len(received), sum(len(chunk) for chunk in received)) + served: Final = _served(wire) + assert _routes(served) == (("GET", f"/v1/files/{file_id}/content"),), served + assert served[0].headers["authorization"] == PROVIDER_AUTHORIZATION, served[0].headers + + +def test_r13_batches_create_retrieve_list_and_cancel_through_the_async_openai_sdk(gateway: Gateway) -> None: + file_id: Final = _identity("file") + batch_id: Final = _identity("batch") + with gateway.scenario() as scenario, wire_server(_batches_peer(file_id, batch_id)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + uploaded: Final = _async_after_worker_sync( + gateway, + lambda client: client.files.create( + file=("batch.jsonl", _batch_line("r1")), purpose="batch", extra_body={"model": model} + ), + ) + created: Final = _async_after_worker_sync( + gateway, + lambda client: client.batches.create( + input_file_id=uploaded.id, + endpoint="/v1/chat/completions", + completion_window="24h", + extra_body={"model": model}, + ), + ) + retrieved: Final = _async_after_worker_sync(gateway, lambda client: client.batches.retrieve(created.id)) + listed: Final = _async_after_worker_sync( + gateway, lambda client: client.batches.list(extra_query={"model": model}) + ) + cancelled: Final = _async_after_worker_sync(gateway, lambda client: client.batches.cancel(created.id)) + listed_ids: Final = tuple(batch.id for batch in listed.data) + assert _upstream_id(uploaded.id, "file-") == file_id, uploaded + assert _upstream_id(created.id, "batch_") == batch_id, created + assert _upstream_id(retrieved.id, "batch_") == batch_id, retrieved + assert batch_id in [_upstream_id(listed_id, "batch_") for listed_id in listed_ids], listed_ids + assert _upstream_id(cancelled.id, "batch_") == batch_id, cancelled + assert cancelled.status == "cancelling", cancelled + served: Final = _served(wire) + assert _routes(served) == ( + ("POST", "/v1/files"), + ("POST", "/v1/batches"), + ("GET", f"/v1/batches/{batch_id}"), + ("POST", f"/v1/batches/{batch_id}/cancel"), + ), served + assert _body(served[1])["input_file_id"] == file_id, served[1].body + assert {request.headers["authorization"] for request in served} == {PROVIDER_AUTHORIZATION}, served + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_r14_fine_tuning_jobs_create_list_retrieve_and_cancel_through_the_sync_openai_sdk( + configured_proxy: ConfiguredProxy, +) -> None: + suffix: Final = uuid.uuid4().hex + job_id: Final = f"ftjob-{suffix}" + with _openai(configured_proxy.gateway) as client: + created: Final = client.fine_tuning.jobs.create( + model=UPSTREAM_MODEL, training_file=f"file-{suffix}", extra_body={"custom_llm_provider": "openai"} + ) + listed: Final = client.fine_tuning.jobs.list(extra_query={"custom_llm_provider": "openai"}) + retrieved: Final = client.fine_tuning.jobs.retrieve(created.id, extra_query={"custom_llm_provider": "openai"}) + cancelled: Final = client.fine_tuning.jobs.cancel(created.id, extra_body={"custom_llm_provider": "openai"}) + assert created.id == job_id, created + assert created.status == "queued", created + assert [job.id for job in listed.data] == [LISTED_JOB], listed + assert retrieved.id == job_id, retrieved + assert retrieved.training_file == f"file-{suffix}", retrieved + assert cancelled.id == job_id and cancelled.status == "cancelled", cancelled + served: Final = _served(configured_proxy.wire) + assert _routes(served) == ( + ("POST", "/v1/fine_tuning/jobs"), + ("GET", "/v1/fine_tuning/jobs"), + ("GET", f"/v1/fine_tuning/jobs/{job_id}"), + ("POST", f"/v1/fine_tuning/jobs/{job_id}/cancel"), + ), served + assert _body(served[0])["training_file"] == f"file-{suffix}", served[0].body + assert {request.headers["authorization"] for request in served} == {PROVIDER_AUTHORIZATION}, served + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_r15_assistants_list_through_raw_httpx(configured_proxy: ConfiguredProxy) -> None: + response: Final = configured_proxy.gateway.request("GET", "/v1/assistants") + assert response.status_code == 200, response.text + body: Final = JSON_OBJECT.validate_json(response.content) + assert body["object"] == "list", response.text + listed: Final = body["data"] + assert isinstance(listed, list), response.text + assert [string_value(object_value(entry)["id"]) for entry in listed] == [ASSISTANT], response.text + served: Final = _served(configured_proxy.wire) + assert _routes(served) == (("GET", "/v1/assistants"),), served + assert served[0].headers["authorization"] == PROVIDER_AUTHORIZATION, served[0].headers + + +def test_r17_chat_completions_on_an_azure_cloudflare_gateway_deployment_through_the_sync_openai_sdk( + gateway: Gateway, +) -> None: + identity: Final = _identity("chatcmpl") + deployment: Final = f"gpt-4o-mini-{uuid.uuid4().hex[:8]}" + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model( + model=f"azure/{deployment}", api_base=f"{wire.url}{CLOUDFLARE_GATEWAY}", api_version=AZURE_API_VERSION + ) + with _openai(gateway) as client: + response: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": TEXT}], extra_body=dict(NO_CACHE) + ) + streamed: Final = tuple( + client.chat.completions.create( + model=model, messages=[{"role": "user", "content": TEXT}], stream=True, extra_body=dict(NO_CACHE) + ) + ) + assert response.id == identity, response + assert response.choices[0].message.content == TEXT, response + assert {chunk.id for chunk in streamed} == {identity}, streamed + assert "".join(_chunk_text(chunk) for chunk in streamed) == TEXT, streamed + served: Final = _served(wire) + assert _routes(served) == (("POST", f"{CLOUDFLARE_GATEWAY}/{deployment}/chat/completions"),) * 2, served + assert {request.target.split("?", 1)[1] for request in served} == {f"api-version={AZURE_API_VERSION}"}, served + assert {request.headers["api-key"] for request in served} == {PROVIDER_KEY}, served + assert [_streams(request) for request in served] == [False, True], served + assert _spend_row(identity)["model"] == f"azure/{deployment}" + + +def test_r18_audio_speech_answers_with_the_upstream_audio_type_through_the_sync_openai_sdk(gateway: Gateway) -> None: + with gateway.scenario() as scenario, wire_server(_speech_peer("audio/ogg")) as wire: + model: Final = scenario.model(model=f"openai/{SPEECH_UPSTREAM_MODEL}", api_base=f"{wire.url}/v1") + speech: Final = _speech_after_worker_sync(gateway, model) + assert not isinstance(speech, openai.APIStatusError), speech + assert speech.status_code == 200, speech.text + assert speech.headers["content-type"] == "audio/ogg", dict(speech.headers) + assert speech.content == SPEECH_AUDIO, speech.content[:16] + served: Final = _served(wire) + assert _routes(served) == (("POST", "/v1/audio/speech"),), served + assert served[0].headers["authorization"] == PROVIDER_AUTHORIZATION, served[0].headers + upstream_body: Final = _body(served[0]) + assert (upstream_body["model"], upstream_body["input"], upstream_body["response_format"]) == ( + SPEECH_UPSTREAM_MODEL, + TEXT, + "mp3", + ), upstream_body + assert _spend_row(speech.headers["x-litellm-call-id"])["call_type"] == "aspeech" + + +def test_s01_chat_completions_upstream_400_reaches_the_sync_openai_sdk_as_a_bad_request(gateway: Gateway) -> None: + with gateway.scenario() as scenario, wire_server(_status_peer(openai_error(400))) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1", max_retries=0) + with _openai(gateway) as client, pytest.raises(openai.BadRequestError) as raised: + client.chat.completions.create(model=model, messages=[_user_message()], extra_body=dict(NO_CACHE)) + assert raised.value.status_code == 400, raised.value + assert "scripted 400" in raised.value.response.text, raised.value.response.text + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) + + +def test_s02_chat_completions_upstream_429_keeps_retry_after_through_raw_httpx(gateway: Gateway) -> None: + limited: Final = Reply(status=429, body=openai_error(429).body, headers={"retry-after": "7"}) + with gateway.scenario() as scenario, wire_server(_status_peer(limited)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1", max_retries=0) + response: Final = _chat(gateway, model) + assert response.status_code == 429, response.text + assert response.headers.get("llm_provider-retry-after") == "7", dict(response.headers) + assert "scripted 429" in _error_message(response), response.text + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) + + +def test_s03_chat_completions_upstream_500_reaches_the_sync_openai_sdk_as_a_server_error(gateway: Gateway) -> None: + with gateway.scenario() as scenario, wire_server(_status_peer(openai_error(500))) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1", max_retries=0) + with _openai(gateway) as client, pytest.raises(openai.InternalServerError) as raised: + client.chat.completions.create(model=model, messages=[_user_message()], extra_body=dict(NO_CACHE)) + assert raised.value.status_code == 500, raised.value + assert "scripted 500" in raised.value.response.text, raised.value.response.text + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) + + +def test_s04_chat_completions_body_timeout_shapes_answer_without_breaking_the_proxy(gateway: Gateway) -> None: + identity: Final = _identity("chatcmpl") + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + for shape, timeout in REJECTED_CHAT_TIMEOUTS.items(): + _assert_rejected(_chat(gateway, model, {"timeout": timeout}), shape) + for shape, timeout in ACCEPTED_CHAT_TIMEOUTS.items(): + _assert_served(_chat(gateway, model, {"timeout": timeout}), identity, shape) + duplicated: Final = gateway.client.post( + "/v1/chat/completions", + content=_duplicated_timeout_body(model), + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + _assert_served(duplicated, identity, "duplicated key") + _assert_served(_chat(gateway, model), identity, "follow-up") + assert len(_served(wire)) == len(ACCEPTED_CHAT_TIMEOUTS) + 2 + + +def test_s05_model_new_timeout_shapes_are_rejected_or_stored_without_breaking_the_proxy(gateway: Gateway) -> None: + identity: Final = _identity("chatcmpl") + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + for shape, timeout in REJECTED_DEPLOYMENT_TIMEOUTS.items(): + _assert_rejected(_new_model(gateway, f"{wire.url}/v1", timeout), shape) + for shape, timeout in ACCEPTED_DEPLOYMENT_TIMEOUTS.items(): + _register_accepted_model(scenario, _new_model(gateway, f"{wire.url}/v1", timeout), shape) + _assert_served(_chat(gateway, model), identity, "follow-up") + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) + + +def test_s06_upstream_errors_on_completions_moderations_and_files_reach_the_sync_openai_sdk( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + with wire_server(_status_peer(openai_error(401))) as unauthorized: + model: Final = scenario.model( + model=INSTRUCT_DEPLOYMENT_MODEL, api_base=f"{unauthorized.url}/v1", max_retries=0 + ) + with _openai(gateway) as client, pytest.raises(openai.AuthenticationError) as completion_error: + client.completions.create(model=model, prompt=TEXT, extra_body=dict(NO_CACHE)) + assert completion_error.value.status_code == 401, completion_error.value + assert "scripted 401" in completion_error.value.response.text, completion_error.value.response.text + assert _routes(_served(unauthorized)) == (("POST", "/v1/completions"),) + with wire_server(_status_peer(openai_error(400))) as rejecting: + rejected_model: Final = scenario.model(api_base=f"{rejecting.url}/v1", max_retries=0) + moderation_error: Final = _moderation_after_worker_sync(gateway, rejected_model) + assert isinstance(moderation_error, openai.BadRequestError), moderation_error + assert "scripted 400" in moderation_error.response.text, moderation_error.response.text + with _openai(gateway) as client: + file_error: Final = _outcome_after_worker_sync( + lambda: client.files.create( + file=("batch.jsonl", _batch_line("rejected")), + purpose="batch", + extra_body={"model": rejected_model}, + ) + ) + assert isinstance(file_error, openai.BadRequestError), file_error + assert "scripted 400" in file_error.response.text, file_error.response.text + assert _routes(_served(rejecting)) == (("POST", "/v1/moderations"), ("POST", "/v1/files")) + + +@pytest.mark.parametrize("provider", tuple(PAYMENT_REQUIRED_DEPLOYMENTS)) +def test_s07_chat_completions_upstream_402_reaches_the_sync_openai_sdk_as_payment_required( + gateway: Gateway, provider: str +) -> None: + with gateway.scenario() as scenario, wire_server(_status_peer(openai_error(402))) as wire: + model: Final = scenario.model(model=PAYMENT_REQUIRED_DEPLOYMENTS[provider], api_base=wire.url, max_retries=0) + with _openai(gateway) as client, pytest.raises(openai.APIStatusError) as raised: + client.chat.completions.create(model=model, messages=[_user_message()], extra_body=dict(NO_CACHE)) + assert raised.value.status_code == 402, raised.value + assert "scripted 402" in raised.value.response.text, raised.value.response.text + assert len(_served(wire)) == 1 + + +def test_e01_deployment_timeout_answers_408_and_a_body_timeout_on_the_same_deployment_recovers( + gateway: Gateway, +) -> None: + identity: Final = _identity("chatcmpl") + gate: Final = threading.Event() + with gateway.scenario() as scenario, wire_server(_held_chat_peer(gate, identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1", timeout=0.5, max_retries=0) + timed_out: Final = _chat(gateway, model) + gate.set() + assert timed_out.status_code == 408, timed_out.text + assert "timeout" in _error_message(timed_out).lower(), timed_out.text + recovered: Final = _chat(gateway, model, {"timeout": 5}) + assert _response_id(recovered) == identity, recovered.text + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) * 2 + + +def test_e02_three_identical_no_cache_requests_through_raw_httpx_land_three_spend_rows(gateway: Gateway) -> None: + served_ids: Final[SimpleQueue[str]] = SimpleQueue() + with gateway.scenario() as scenario, wire_server(_fresh_chat_peer(served_ids)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + received: Final = ( + _response_id(_chat(gateway, model)), + _response_id(_chat(gateway, model)), + _response_id(_chat(gateway, model)), + ) + assert len(set(received)) == 3, received + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) * 3 + assert {served_ids.get_nowait() for _ in range(3)} == set(received), received + assert _spend_request_ids(received) == frozenset(received) + + +def test_e03_deployment_timeout_as_a_string_still_serves_through_raw_httpx(gateway: Gateway) -> None: + identity: Final = _identity("chatcmpl") + with gateway.scenario() as scenario, wire_server(_chat_peer(identity)) as wire: + model: Final = scenario.model(api_base=f"{wire.url}/v1", timeout="30") + assert _response_id(_chat(gateway, model)) == identity + assert _routes(_served(wire)) == (("POST", "/v1/chat/completions"),) + assert _spend_row(identity)["model"] == DEPLOYMENT_MODEL diff --git a/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py b/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py new file mode 100644 index 00000000000..0bcc9fc4340 --- /dev/null +++ b/tests/integration/compatibility/test_openai_sdk_client_paths_chaos.py @@ -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} " + + +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)), + ) diff --git a/tests/integration/providers/test_databricks_decisions_wire.py b/tests/integration/providers/test_databricks_decisions_wire.py index 5c3b7d1b1a1..93dc2ea04fa 100644 --- a/tests/integration/providers/test_databricks_decisions_wire.py +++ b/tests/integration/providers/test_databricks_decisions_wire.py @@ -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 diff --git a/tests/integration/providers/test_decisions_openai_format_wire.py b/tests/integration/providers/test_decisions_openai_format_wire.py new file mode 100644 index 00000000000..f0e7163e264 --- /dev/null +++ b/tests/integration/providers/test_decisions_openai_format_wire.py @@ -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() diff --git a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py index fd8879e02d4..e39f7e360f6 100644 --- a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py +++ b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py @@ -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)) diff --git a/tests/integration/providers/test_realtime_unguarded_voice_turns.py b/tests/integration/providers/test_realtime_unguarded_voice_turns.py new file mode 100644 index 00000000000..396b266f45f --- /dev/null +++ b/tests/integration/providers/test_realtime_unguarded_voice_turns.py @@ -0,0 +1,1004 @@ +"""Without a transcript guardrail the upstream receives exactly the client's realtime frames. + +The proxy sends its own ``response.create`` after a voice transcript only when a ``realtime_input_transcription`` +guardrail turned the provider's VAD auto-response off. Without one the client owns its turns: one +``response.create`` per turn reaches the upstream, carrying the client's own metadata, and no bare one from the +proxy. The scripted upstream answers one event per ``input_audio_buffer.commit`` or ``response.create``, so a +proxy-sent ``response.create`` shows up both as an extra frame and as a reply the client never asked for. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import os +import uuid +from collections import Counter +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from dataclasses import dataclass +from hashlib import sha256 +from itertools import chain +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import websockets +from openai import AsyncOpenAI, OpenAI +from openai.resources.realtime.realtime import AsyncRealtimeConnection, RealtimeConnection +from openai.types.realtime import RealtimeSessionCreateRequestParam +from pydantic import JsonValue +from websockets.asyncio.client import ClientConnection +from websockets.exceptions import ConnectionClosed + +from tests.integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from tests.integration._support.database import read_rows +from tests.integration._support.process import ( + OwnedProxy, + UpstreamSlot, + graceful_stop_seconds, + owned_proxy_process, + owned_upstream, +) +from tests.integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario +from tests.integration.cost_calculation.cost_tracking_case import RealtimeResponse + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +WORKERS: Final = 2 +BURST: Final = 20 +OPEN_SESSIONS: Final = 6 +MAX_OPEN_SESSIONS: Final = 24 +VOICE_MODEL: Final = "gpt-realtime-2.1" +XAI_VOICE_MODEL: Final = "grok-voice" +LIVE_MODEL: Final = "gemini-3.8-live" +LIVE_LOCATION: Final = "us-central1" +AZURE_API_VERSION: Final = "2025-04-01-preview" +TRANSCRIPT: Final = "what is the weather like today" +REALTIME_HOOK: Final = "realtime_input_transcription" +PROMPT_GUARDRAIL: Final = "prompt-filter" +PROMPT_WORD: Final = "pineapple" +MIXED_GUARDRAIL: Final = "prompt-and-transcript-filter" +MIXED_WORD: Final = "anchovy" +TRANSCRIPT_COMPLETED: Final = "conversation.item.input_audio_transcription.completed" +RESPONSE_DONE: Final = "response.done" +HANDSHAKE_REFUSED: Final = "Upstream realtime handshake rejected with HTTP 403" +REFUSAL_CLOSE: Final = 1008 +CREATED_SECONDS: Final = 60.0 +STEP_SECONDS: Final = 30.0 +FIVE_KB: Final = "x" * 5120 +PCM_200MS_24KHZ: Final = base64.b64encode(bytes(9600)).decode() +SPEND_SQL: Final = 'SELECT call_type FROM "LiteLLM_SpendLogs" WHERE api_key = %s' +ROUTES: Final = ("/v1/realtime", "/realtime", "/openai/v1/realtime") +BETA_HEADERS: Final = {"OpenAI-Beta": "realtime=v1"} +COMMIT: Final[dict[str, JsonValue]] = {"type": "input_audio_buffer.commit"} +APPEND: Final[dict[str, JsonValue]] = {"type": "input_audio_buffer.append", "audio": PCM_200MS_24KHZ} +SDK_PUSH_TO_TALK: Final[RealtimeSessionCreateRequestParam] = { + "type": "realtime", + "audio": {"input": {"turn_detection": None}}, +} +PUSH_TO_TALK: Final = JSON_OBJECT.validate_python(SDK_PUSH_TO_TALK) +BETA_PUSH_TO_TALK: Final[dict[str, JsonValue]] = {"turn_detection": None} +VOICE_UPDATES: Final = ( + pytest.param({"type": "realtime", "instructions": "answer briefly"}, id="turn-detection-unset"), + pytest.param( + {"type": "realtime", "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}}, + id="server-vad-without-auto-response", + ), + pytest.param( + { + "type": "realtime", + "audio": {"input": {"turn_detection": {"type": "semantic_vad", "create_response": False}}}, + }, + id="semantic-vad-without-auto-response", + ), +) +TRANSCRIPT_FIELDS: Final = ( + pytest.param({}, id="missing"), + pytest.param({"transcript": ""}, id="empty"), + pytest.param({"transcript": FIVE_KB}, id="five-kilobytes"), + pytest.param({"transcript": None}, id="null"), + pytest.param({"transcript": 123}, id="int"), + pytest.param({"transcript": ["a"]}, id="list"), + pytest.param({"transcript": {"text": "a"}}, id="object"), +) +VOICE_UPDATE: Final[dict[str, JsonValue]] = {"type": "realtime", "instructions": "answer briefly"} +VOICE_UPDATE_FORWARDED: Final[dict[str, JsonValue]] = { + "type": "realtime", + "instructions": "answer briefly", + "audio": {"input": {"turn_detection": {"create_response": False}}}, +} +TURN_DETECTION_TYPE_MISSING: Final = "Missing required parameter: 'session.audio.input.turn_detection.type'." +INJECTED_UPDATE: Final[dict[str, JsonValue]] = { + "type": "realtime", + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, +} +PROXY_RESPONSE_CREATE: Final[dict[str, JsonValue]] = {"type": "response.create"} + + +@dataclass(frozen=True, slots=True) +class Session: + events: tuple[dict[str, JsonValue], ...] + + @property + def types(self) -> tuple[str, ...]: + return tuple(string_value(event["type"]) for event in self.events) + + @property + def transcript_fields(self) -> tuple[dict[str, JsonValue], ...]: + completed: Final = tuple(event for event in self.events if event["type"] == TRANSCRIPT_COMPLETED) + return tuple({key: value for key, value in event.items() if key == "transcript"} for event in completed) + + @property + def response_ids(self) -> tuple[str, ...]: + done: Final = tuple(event for event in self.events if event["type"] == RESPONSE_DONE) + return tuple(string_value(object_value(event["response"])["id"]) for event in done) + + @property + def error_messages(self) -> tuple[str, ...]: + errors: Final = tuple(event for event in self.events if event["type"] == "error") + return tuple(string_value(object_value(event["error"])["message"]) for event in errors) + + @property + def close_code(self) -> JsonValue: + closes: Final = tuple(event for event in self.events if event["type"] == "closed") + return closes[-1]["code"] if closes else None + + +@dataclass(frozen=True, slots=True) +class Step: + frames: tuple[dict[str, JsonValue], ...] + until: str + + +@dataclass(frozen=True, slots=True) +class Outage: + turns: tuple[Session, ...] + observed: tuple[dict[str, JsonValue], ...] + held: tuple[Session, ...] + + +def _ws_base(http_url: str) -> str: + return http_url.replace("https://", "wss://").replace("http://", "ws://") + + +def _gateway_url(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def _owned_url(owned: OwnedProxy) -> str: + return _gateway_url(owned.gateway) + + +def _transcript_event(fields: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + "type": TRANSCRIPT_COMPLETED, + "event_id": "evt_$UNIQUE_ID", + "item_id": "item_$UNIQUE_ID", + "content_index": 0, + **fields, + "usage": {"type": "duration", "seconds": 1.5}, + } + + +def _done_event() -> dict[str, JsonValue]: + return { + "type": RESPONSE_DONE, + "event_id": "evt_$UNIQUE_ID", + "response": { + "id": "resp_$UNIQUE_ID", + "object": "realtime.response", + "status": "completed", + "output": [], + "usage": {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}, + }, + } + + +def _turns(fields: tuple[Mapping[str, JsonValue], ...]) -> RealtimeResponse: + events: Final = tuple(chain.from_iterable((_transcript_event(turn), _done_event()) for turn in fields)) + return RealtimeResponse(content_type="application/x-realtime", events=events) + + +def _spoken(transcript: str = TRANSCRIPT, *, turns: int = 1) -> RealtimeResponse: + return _turns(({"transcript": transcript},) * turns) + + +def _live_turn() -> RealtimeResponse: + return RealtimeResponse( + content_type="application/x-realtime", + events=( + {"serverContent": {"inputTranscription": {"text": TRANSCRIPT}}}, + {"serverContent": {"modelTurn": {"parts": [{"text": "scripted $REQUEST_ID"}]}}}, + { + "serverContent": {"turnComplete": True}, + "usageMetadata": {"promptTokenCount": 7, "responseTokenCount": 5, "totalTokenCount": 12}, + }, + ), + ) + + +def _scripted(scenario: Scenario, response: RealtimeResponse, *, control_url: str | None = None) -> ScenarioHandle: + scenario_id: Final = f"realtime-unguarded-{uuid.uuid4().hex[:12]}" + handle: Final = ( + register_scenario(scenario_id, response) + if control_url is None + else register_scenario(scenario_id, response, control_url=control_url) + ) + scenario.cleanups.callback(delete_scenario, handle) + return handle + + +def _openai_deployment(scenario: Scenario, scenario_id: str) -> str: + return scenario.model(model=f"openai/{VOICE_MODEL}", api_key=scenario_id) + + +def _named_deployment(creator: Gateway, scenario: Scenario, scenario_id: str) -> str: + name: Final = f"realtime-unguarded-owned-{uuid.uuid4().hex[:8]}" + created: Final = creator.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": f"openai/{VOICE_MODEL}", + "api_key": scenario_id, + "api_base": f"{creator.upstream_url}/v1", + }, + "model_info": {}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return name + + +def _live_credentials(upstream_url: str) -> str: + return json.dumps( + { + "type": "external_account", + "audience": "synthetic-vertex-audience", + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "token_url": f"{upstream_url}/_oauth/token", + "credential_source": {"url": f"{upstream_url}/health"}, + } + ) + + +def _live_deployment(gateway: Gateway, scenario: Scenario, project: str) -> str: + name: Final = f"vertex-live-unguarded-{uuid.uuid4().hex[:12]}" + upstream: Final = gateway.upstream_url.rstrip("/") + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": f"vertex_ai/{LIVE_MODEL}", + "vertex_project": project, + "vertex_credentials": _live_credentials(upstream), + "vertex_location": LIVE_LOCATION, + "api_base": upstream, + }, + "model_info": {"mode": "realtime"}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return name + + +def _client_response_create(label: str) -> dict[str, JsonValue]: + return {"type": "response.create", "response": {"metadata": {"turn": label}}} + + +def _update_frame(update: dict[str, JsonValue]) -> dict[str, JsonValue]: + return {"type": "session.update", "session": update} + + +def _turn(label: str) -> tuple[Step, Step]: + return Step((COMMIT,), TRANSCRIPT_COMPLETED), Step((_client_response_create(label),), RESPONSE_DONE) + + +def _client_turns(update: dict[str, JsonValue], labels: tuple[str, ...]) -> tuple[Step, ...]: + return (Step((_update_frame(update),), "session.updated"), *chain.from_iterable(map(_turn, labels))) + + +def _client_frames(update: dict[str, JsonValue], labels: tuple[str, ...]) -> tuple[dict[str, JsonValue], ...]: + turns: Final = chain.from_iterable((COMMIT, _client_response_create(label)) for label in labels) + return (_update_frame(update), *turns) + + +def _close_code(closed: ConnectionClosed) -> int: + return 1006 if closed.rcvd is None else closed.rcvd.code + + +async def _next_event(socket: ClientConnection, deadline: float) -> dict[str, JsonValue] | None: + remaining: Final = deadline - asyncio.get_running_loop().time() + try: + return JSON_OBJECT.validate_json(await asyncio.wait_for(socket.recv(), max(remaining, 0.01))) + except TimeoutError: + return None + + +async def _until(socket: ClientConnection, until: str) -> AsyncIterator[dict[str, JsonValue]]: + deadline: Final = asyncio.get_running_loop().time() + STEP_SECONDS + while (event := await _next_event(socket, deadline)) is not None: + yield event + if event.get("type") == until: + return + yield {"type": "timeout"} + + +async def _steps_after_created( + socket: ClientConnection, steps: tuple[Step, ...] +) -> AsyncIterator[dict[str, JsonValue]]: + for step in steps: + for frame in step.frames: + await socket.send(json.dumps(frame)) + async for event in _until(socket, step.until): + yield event + if event["type"] == "timeout": + return + + +async def _stepped(socket: ClientConnection, steps: tuple[Step, ...]) -> AsyncIterator[dict[str, JsonValue]]: + try: + first: Final = JSON_OBJECT.validate_json(await asyncio.wait_for(socket.recv(), CREATED_SECONDS)) + yield first + if first.get("type") != "session.created": + async for message in socket: + yield JSON_OBJECT.validate_json(message) + return + async for event in _steps_after_created(socket, steps): + yield event + except ConnectionClosed as closed: + yield {"type": "closed", "code": _close_code(closed)} + + +def _headers(key: str, extra: Mapping[str, str] | None) -> dict[str, str]: + return {"Authorization": f"Bearer {key}", **(extra or {})} + + +async def _driven(url: str, key: str, steps: tuple[Step, ...], headers: Mapping[str, str] | None) -> Session: + async with websockets.connect(url, additional_headers=_headers(key, headers)) as socket: + return Session(tuple([event async for event in _stepped(socket, steps)])) + + +def _session_url(proxy_url: str, query: str, path: str = "/v1/realtime") -> str: + return f"{_ws_base(proxy_url)}{path}?{query}" + + +def _drive( + proxy_url: str, + query: str, + key: str, + steps: tuple[Step, ...], + *, + path: str = "/v1/realtime", + headers: Mapping[str, str] | None = None, +) -> Session: + return asyncio.run(_driven(_session_url(proxy_url, query, path), key, steps, headers)) + + +def _observed_all(upstream_url: str) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client( + base_url=upstream_url, timeout=5, trust_env=False, verify=not upstream_url.startswith("https://") + ) as upstream: + observations: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content) + requests: Final = observations["requests"] + assert isinstance(requests, list), observations + return tuple(map(object_value, requests)) + + +def _belongs(request: dict[str, JsonValue], scenario_id: str) -> bool: + return request["authorization"] == f"Bearer {scenario_id}" or request["api_key"] == scenario_id + + +def _observed(upstream_url: str, scenario_id: str) -> tuple[dict[str, JsonValue], ...]: + return tuple(request for request in _observed_all(upstream_url) if _belongs(request, scenario_id)) + + +def _upgrades(observed: tuple[dict[str, JsonValue], ...]) -> tuple[tuple[JsonValue, JsonValue], ...]: + upgrades: Final = tuple(request for request in observed if request["method"] == "WEBSOCKET") + return tuple((request["path"], object_value(request["body"]).get("query")) for request in upgrades) + + +def _sent(observed: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return tuple(object_value(request["body"]) for request in observed if request["method"] == "WEBSOCKET_FRAME") + + +def _frame_counts(frames: Iterator[dict[str, JsonValue]] | tuple[dict[str, JsonValue], ...]) -> Counter[str]: + return Counter(json.dumps(frame, sort_keys=True) for frame in frames) + + +def _spend_rows(key: str, count: int) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows(SPEND_SQL, (sha256(key.encode()).hexdigest(),)), + lambda rows: len(rows) == count, + seconds=70, + ) + + +def _assert_billed(key: str, sessions: int) -> None: + rows: Final = _spend_rows(key, sessions) + assert [row["call_type"] for row in rows] == ["_arealtime"] * sessions, rows + + +def _assert_turns_answered(session: Session, scenario_id: str, fields: tuple[Mapping[str, JsonValue], ...]) -> None: + assert session.types == ( + "session.created", + "session.updated", + *(TRANSCRIPT_COMPLETED, RESPONSE_DONE) * len(fields), + ), session + assert session.transcript_fields == fields, session + assert session.error_messages == (), session + assert len(set(session.response_ids)) == len(fields), session + assert all(response_id.startswith(f"resp_{scenario_id}-") for response_id in session.response_ids), session + + +def _assert_client_turns( + session: Session, + observed: tuple[dict[str, JsonValue], ...], + scenario_id: str, + update: dict[str, JsonValue], + labels: tuple[str, ...], + *, + fields: tuple[Mapping[str, JsonValue], ...] | None = None, +) -> None: + _assert_turns_answered(session, scenario_id, fields or ({"transcript": TRANSCRIPT},) * len(labels)) + assert _sent(observed) == _client_frames(update, labels), _sent(observed) + + +def _content_filter(name: str, mode: JsonValue, *, default_on: bool, word: str) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": mode, + "default_on": default_on, + "blocked_words": [{"keyword": word, "action": "BLOCK"}], + }, + } + + +def _write_config(directory: Path, ca_bundle: Path) -> Path: + config: Final = directory / f"realtime_unguarded_{uuid.uuid4().hex[:8]}.yaml" + config.write_text( + json.dumps( + { + "guardrails": [ + _content_filter(PROMPT_GUARDRAIL, "pre_call", default_on=True, word=PROMPT_WORD), + _content_filter(MIXED_GUARDRAIL, ["pre_call", REALTIME_HOOK], default_on=False, word=MIXED_WORD), + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "store_model_in_db": True, + "proxy_batch_write_at": 1, + "proxy_batch_polling_interval": 1, + }, + "litellm_settings": { + "ssl_verify": str(ca_bundle), + "enable_redis_auth_cache": True, + "cache": True, + "cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"}, + }, + "router_settings": {"disable_cooldowns": True}, + } + ) + ) + return config + + +def _ca_bundle(slot: UpstreamSlot) -> Path: + assert slot.certificate is not None, "the TLS upstream carries its own certificate" + return slot.certificate.certificate + + +@pytest.fixture(scope="module") +def tls_slot(tmp_path_factory: pytest.TempPathFactory) -> Iterator[UpstreamSlot]: + with owned_upstream(tmp_path_factory.mktemp("unguarded-tls-upstream"), tls=True) as slot: + yield slot + + +@pytest.fixture(scope="module") +def tls_upstream(tls_slot: UpstreamSlot) -> str: + return tls_slot.url + + +@pytest.fixture(scope="module") +def guard_proxy(tmp_path_factory: pytest.TempPathFactory, tls_slot: UpstreamSlot) -> Iterator[OwnedProxy]: + directory: Final = tmp_path_factory.mktemp("realtime-unguarded-proxy") + ca_bundle: Final = _ca_bundle(tls_slot) + with ( + gateway_from_environment() as rig, + owned_proxy_process( + rig, + directory, + {"DATABASE_URL": os.environ["DATABASE_URL"], "SSL_VERIFY": str(ca_bundle)}, + config=_write_config(directory, ca_bundle), + workers=WORKERS, + ) as owned, + ): + yield owned + + +@pytest.mark.parametrize("path", ROUTES) +def test_push_to_talk_turn_reaches_the_upstream_as_the_clients_own_frames(gateway: Gateway, path: str) -> None: + with gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _spoken()) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id) + session: Final = _drive( + _gateway_url(gateway), f"model={model}", key, _client_turns(PUSH_TO_TALK, ("1",)), path=path + ) + observed: Final = _observed(gateway.upstream_url, handle.scenario_id) + _assert_client_turns(session, observed, handle.scenario_id, PUSH_TO_TALK, ("1",)) + assert _upgrades(observed) == (("/v1/realtime", [["model", VOICE_MODEL]]),), observed + _assert_billed(key, 1) + + +@pytest.mark.parametrize("update", VOICE_UPDATES) +def test_voice_session_turn_carries_only_the_clients_response_create( + gateway: Gateway, update: dict[str, JsonValue] +) -> None: + with gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _spoken()) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id) + session: Final = _drive(_gateway_url(gateway), f"model={model}", key, _client_turns(update, ("1",))) + observed: Final = _observed(gateway.upstream_url, handle.scenario_id) + _assert_client_turns(session, observed, handle.scenario_id, update, ("1",)) + _assert_billed(key, 1) + + +def test_beta_protocol_push_to_talk_turn_reaches_the_upstream_as_the_clients_own_frames(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _spoken()) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id) + session: Final = _drive( + _gateway_url(gateway), + f"model={model}", + key, + _client_turns(BETA_PUSH_TO_TALK, ("1",)), + headers=BETA_HEADERS, + ) + observed: Final = _observed(gateway.upstream_url, handle.scenario_id) + _assert_client_turns(session, observed, handle.scenario_id, BETA_PUSH_TO_TALK, ("1",)) + _assert_billed(key, 1) + + +def test_two_turns_with_the_same_transcript_get_one_reply_each(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _spoken(turns=2)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id) + session: Final = _drive(_gateway_url(gateway), f"model={model}", key, _client_turns(PUSH_TO_TALK, ("1", "2"))) + observed: Final = _observed(gateway.upstream_url, handle.scenario_id) + _assert_client_turns(session, observed, handle.scenario_id, PUSH_TO_TALK, ("1", "2")) + _assert_billed(key, 1) + + +async def _sdk_async_events(connection: AsyncRealtimeConnection, until: str) -> AsyncIterator[dict[str, JsonValue]]: + async for event in connection: + yield JSON_OBJECT.validate_python(event.model_dump()) + if event.type == until: + return + + +def _sdk_sync_events(connection: RealtimeConnection, until: str) -> Iterator[dict[str, JsonValue]]: + for event in connection: + yield JSON_OBJECT.validate_python(event.model_dump()) + if event.type == until: + return + + +async def _sdk_async_turn(proxy_url: str, key: str, model: str) -> Session: + client: Final = AsyncOpenAI(api_key=key, base_url=f"{proxy_url}/v1", websocket_base_url=f"{_ws_base(proxy_url)}/v1") + async with client.realtime.connect(model=model) as connection: + created: Final = [event async for event in _sdk_async_events(connection, "session.created")] + await connection.session.update(session=SDK_PUSH_TO_TALK) + updated: Final = [event async for event in _sdk_async_events(connection, "session.updated")] + await connection.input_audio_buffer.commit() + transcribed: Final = [event async for event in _sdk_async_events(connection, TRANSCRIPT_COMPLETED)] + await connection.response.create(response={"metadata": {"turn": "1"}}) + answered: Final = [event async for event in _sdk_async_events(connection, RESPONSE_DONE)] + return Session((*created, *updated, *transcribed, *answered)) + + +def _sdk_sync_turn(proxy_url: str, key: str, model: str) -> Session: + client: Final = OpenAI(api_key=key, base_url=f"{proxy_url}/v1", websocket_base_url=f"{_ws_base(proxy_url)}/v1") + with client.realtime.connect(model=model) as connection: + created: Final = tuple(_sdk_sync_events(connection, "session.created")) + connection.session.update(session=SDK_PUSH_TO_TALK) + updated: Final = tuple(_sdk_sync_events(connection, "session.updated")) + connection.input_audio_buffer.commit() + transcribed: Final = tuple(_sdk_sync_events(connection, TRANSCRIPT_COMPLETED)) + connection.response.create(response={"metadata": {"turn": "1"}}) + answered: Final = tuple(_sdk_sync_events(connection, RESPONSE_DONE)) + return Session((*created, *updated, *transcribed, *answered)) + + +@pytest.mark.parametrize("client", ["async", "sync"]) +def test_openai_sdk_push_to_talk_turn_reaches_the_upstream_as_the_clients_own_frames( + gateway: Gateway, client: str +) -> None: + with gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _spoken()) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id) + session: Final = ( + asyncio.run(_sdk_async_turn(_gateway_url(gateway), key, model)) + if client == "async" + else _sdk_sync_turn(_gateway_url(gateway), key, model) + ) + observed: Final = _observed(gateway.upstream_url, handle.scenario_id) + _assert_client_turns(session, observed, handle.scenario_id, PUSH_TO_TALK, ("1",)) + _assert_billed(key, 1) + + +def test_azure_push_to_talk_turn_reaches_the_upstream_as_the_clients_own_frames( + guard_proxy: OwnedProxy, tls_upstream: str +) -> None: + with guard_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _spoken(), control_url=tls_upstream) + key: Final = scenario.key() + model: Final = scenario.model( + model=f"azure/{VOICE_MODEL}", + api_key=handle.scenario_id, + api_base=_ws_base(tls_upstream), + api_version=AZURE_API_VERSION, + ) + session: Final = _drive(_owned_url(guard_proxy), f"model={model}", key, _client_turns(PUSH_TO_TALK, ("1",))) + observed: Final = _observed(tls_upstream, handle.scenario_id) + _assert_client_turns(session, observed, handle.scenario_id, PUSH_TO_TALK, ("1",)) + assert _upgrades(observed) == (("/openai/v1/realtime", [["model", VOICE_MODEL]]),), observed + _assert_billed(key, 1) + + +def test_xai_push_to_talk_turn_reaches_the_upstream_as_the_clients_own_frames(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _spoken()) + key: Final = scenario.key() + model: Final = scenario.model(model=f"xai/{XAI_VOICE_MODEL}", api_key=handle.scenario_id) + session: Final = _drive(_gateway_url(gateway), f"model={model}", key, _client_turns(PUSH_TO_TALK, ("1",))) + observed: Final = _observed(gateway.upstream_url, handle.scenario_id) + _assert_client_turns(session, observed, handle.scenario_id, PUSH_TO_TALK, ("1",)) + assert _upgrades(observed) == (("/v1/realtime", [["model", XAI_VOICE_MODEL]]),), observed + _assert_billed(key, 1) + + +def test_prompt_guardrail_alone_leaves_the_turn_to_the_client_when_the_transcript_has_its_word( + guard_proxy: OwnedProxy, +) -> None: + transcript: Final = f"please add {PROMPT_WORD} to the order" + with guard_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _spoken(transcript)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id) + session: Final = _drive(_owned_url(guard_proxy), f"model={model}", key, _client_turns(PUSH_TO_TALK, ("1",))) + observed: Final = _observed(guard_proxy.gateway.upstream_url, handle.scenario_id) + _assert_client_turns( + session, observed, handle.scenario_id, PUSH_TO_TALK, ("1",), fields=({"transcript": transcript},) + ) + _assert_billed(key, 1) + + +def test_vertex_live_voice_turn_sends_the_upstream_only_the_setup_and_the_audio(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _live_turn()) + key: Final = scenario.key() + model: Final = _live_deployment(gateway, scenario, handle.scenario_id) + session: Final = _drive(_gateway_url(gateway), f"model={model}", key, (Step((APPEND,), RESPONSE_DONE),)) + observed: Final = _observed(gateway.upstream_url, handle.scenario_id) + assert session.types[0] == "session.created", session + assert session.types[-1] == RESPONSE_DONE, session + assert session.transcript_fields == ({"transcript": TRANSCRIPT},), session + assert session.error_messages == (), session + assert len(session.response_ids) == 1, session + sent: Final = _sent(observed) + assert [sorted(frame) for frame in sent] == [["setup"], ["realtimeInput"]], sent + audio: Final = object_value(object_value(sent[1]["realtimeInput"])["audio"]) + assert audio["data"] == PCM_200MS_24KHZ, audio + _assert_billed(key, 1) + + +@pytest.mark.parametrize("fields", TRANSCRIPT_FIELDS) +def test_transcript_field_shape_is_relayed_and_only_the_clients_response_create_follows( + gateway: Gateway, fields: dict[str, JsonValue] +) -> None: + with gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _turns((fields,))) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id) + session: Final = _drive(_gateway_url(gateway), f"model={model}", key, _client_turns(PUSH_TO_TALK, ("1",))) + observed: Final = _observed(gateway.upstream_url, handle.scenario_id) + _assert_client_turns(session, observed, handle.scenario_id, PUSH_TO_TALK, ("1",), fields=(fields,)) + _assert_billed(key, 1) + + +def test_refused_upstream_upgrade_reaches_the_client_and_the_next_session_keeps_its_own_turns( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + unregistered: Final = f"realtime-unguarded-unknown-{uuid.uuid4().hex[:12]}" + handle: Final = _scripted(scenario, _spoken()) + key: Final = scenario.key() + refused_model: Final = _openai_deployment(scenario, unregistered) + model: Final = _openai_deployment(scenario, handle.scenario_id) + refused: Final = _drive( + _gateway_url(gateway), f"model={refused_model}", key, _client_turns(PUSH_TO_TALK, ("1",)) + ) + assert refused.types == ("error", "closed"), refused + assert refused.error_messages == (HANDSHAKE_REFUSED,), refused + assert refused.close_code == REFUSAL_CLOSE, refused + refused_observed: Final = _observed(gateway.upstream_url, unregistered) + assert _upgrades(refused_observed) == (("/v1/realtime", [["model", VOICE_MODEL]]),), refused_observed + assert _sent(refused_observed) == (), refused_observed + session: Final = _drive(_gateway_url(gateway), f"model={model}", key, _client_turns(PUSH_TO_TALK, ("1",))) + observed: Final = _observed(gateway.upstream_url, handle.scenario_id) + _assert_client_turns(session, observed, handle.scenario_id, PUSH_TO_TALK, ("1",)) + _assert_billed(key, 1) + + +def _labels(count: int) -> tuple[str, ...]: + return tuple(str(index) for index in range(count)) + + +def _every_client_frame(labels: tuple[str, ...]) -> Iterator[dict[str, JsonValue]]: + return chain.from_iterable(_client_frames(PUSH_TO_TALK, (label,)) for label in labels) + + +async def _concurrent_turns(url: str, key: str, labels: tuple[str, ...]) -> tuple[Session, ...]: + sessions: Final = (_driven(url, key, _client_turns(PUSH_TO_TALK, (label,)), None) for label in labels) + return tuple(await asyncio.gather(*sessions)) + + +def test_concurrent_sessions_each_get_only_their_own_response_create(gateway: Gateway) -> None: + labels: Final = _labels(BURST) + with gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _spoken()) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id) + sessions: Final = asyncio.run( + _concurrent_turns(_session_url(_gateway_url(gateway), f"model={model}"), key, labels) + ) + for session in sessions: + _assert_turns_answered(session, handle.scenario_id, ({"transcript": TRANSCRIPT},)) + assert len({response_id for session in sessions for response_id in session.response_ids}) == BURST, sessions + sent: Final = _sent(_observed(gateway.upstream_url, handle.scenario_id)) + assert _frame_counts(sent) == _frame_counts(_every_client_frame(labels)), sent + _assert_billed(key, BURST) + + +async def _frames_until_closed(socket: ClientConnection) -> AsyncIterator[dict[str, JsonValue]]: + try: + async for message in socket: + yield JSON_OBJECT.validate_json(message) + except ConnectionClosed as closed: + yield {"type": "closed", "code": _close_code(closed)} + + +async def _turn_then_hold(url: str, key: str, label: str, turned: asyncio.Queue[Session]) -> Session: + async with websockets.connect(url, additional_headers=_headers(key, None)) as socket: + turn: Final = Session(tuple([event async for event in _stepped(socket, _client_turns(PUSH_TO_TALK, (label,)))])) + await turned.put(turn) + return Session(tuple([frame async for frame in _frames_until_closed(socket)])) + + +async def _drain(turned: asyncio.Queue[Session], count: int) -> tuple[Session, ...]: + return tuple([await turned.get() for _ in range(count)]) + + +async def _burst_through_outage( + url: str, + proxy_url: str, + key: str, + observe: Callable[[], tuple[dict[str, JsonValue], ...]], + stop_upstream: Callable[[], None], +) -> Outage: + turned: Final[asyncio.Queue[Session]] = asyncio.Queue() + holders: Final = tuple(asyncio.ensure_future(_turn_then_hold(url, key, label, turned)) for label in _labels(BURST)) + turns: Final = await asyncio.wait_for(_drain(turned, BURST), 120) + observed: Final = await asyncio.to_thread(observe) + await asyncio.to_thread(stop_upstream) + async with httpx.AsyncClient(base_url=proxy_url, timeout=15, trust_env=False) as client: + liveliness: Final = await client.get("/health/liveliness") + assert liveliness.status_code == 200, liveliness.text + return Outage(turns, observed, tuple(await asyncio.wait_for(asyncio.gather(*holders), 90))) + + +def test_upstream_outage_mid_burst_closes_every_session_after_only_the_clients_frames( + gateway: Gateway, tmp_path: Path +) -> None: + with gateway.scenario() as scenario, owned_upstream(tmp_path) as slot: + scenario_id: Final = f"realtime-unguarded-outage-{uuid.uuid4().hex[:12]}" + register_scenario(scenario_id, _spoken(), control_url=slot.url) + key: Final = scenario.key() + model: Final = scenario.model(model=f"openai/{VOICE_MODEL}", api_key=scenario_id, api_base=slot.url) + proxy_url: Final = _gateway_url(gateway) + outage: Final = asyncio.run( + _burst_through_outage( + _session_url(proxy_url, f"model={model}"), + proxy_url, + key, + lambda: _observed(slot.url, scenario_id), + slot.stop, + ) + ) + for turn in outage.turns: + _assert_turns_answered(turn, scenario_id, ({"transcript": TRANSCRIPT},)) + assert _frame_counts(_sent(outage.observed)) == _frame_counts(_every_client_frame(_labels(BURST))), outage + assert [session.types for session in outage.held] == [("error", "closed")] * BURST, outage.held + slot.start() + register_scenario(scenario_id, _spoken(), control_url=slot.url) + recovered: Final = _drive(proxy_url, f"model={model}", key, _client_turns(PUSH_TO_TALK, ("recovered",))) + _assert_client_turns(recovered, _observed(slot.url, scenario_id), scenario_id, PUSH_TO_TALK, ("recovered",)) + _assert_billed(key, BURST + 1) + + +def _spawned_worker(child: psutil.Process) -> bool: + try: + return any("multiprocessing.spawn" in part for part in child.cmdline()) + except (psutil.NoSuchProcess, psutil.AccessDenied): + return False + + +def _exited(process: psutil.Process) -> bool: + try: + return process.status() == psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return True + + +def _workers(root: psutil.Process) -> tuple[psutil.Process, ...]: + return tuple( + sorted((child for child in root.children() if _spawned_worker(child)), key=lambda process: process.pid) + ) + + +async def _turn_or_close(socket: ClientConnection, label: str) -> Session: + try: + return Session( + tuple([event async for event in _steps_after_created(socket, _client_turns(PUSH_TO_TALK, (label,)))]) + ) + except ConnectionClosed as closed: + return Session(({"type": "closed", "code": _close_code(closed)},)) + + +@dataclass(frozen=True, slots=True) +class WorkerSession: + socket: ClientConnection + pid: int + + +def _accepted_client_ports(worker: psutil.Process, proxy_port: int) -> frozenset[int]: + return frozenset( + connection.raddr.port + for connection in worker.net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.laddr.port == proxy_port and connection.raddr + ) + + +async def _open_session(url: str, key: str, workers: tuple[psutil.Process, ...]) -> WorkerSession: + proxy_port: Final = urlsplit(url).port or 0 + before: Final = tuple(_accepted_client_ports(worker, proxy_port) for worker in workers) + socket: Final = await websockets.connect(url, additional_headers=_headers(key, None)) + created: Final = JSON_OBJECT.validate_json(await asyncio.wait_for(socket.recv(), 60)) + assert created.get("type") == "session.created", created + accepted_by: Final = tuple( + worker.pid for worker, ports in zip(workers, before) if _accepted_client_ports(worker, proxy_port) - ports + ) + assert len(accepted_by) == 1, accepted_by + return WorkerSession(socket, accepted_by[0]) + + +async def _sessions_on_every_worker( + url: str, key: str, workers: tuple[psutil.Process, ...], opened: tuple[WorkerSession, ...] +) -> tuple[WorkerSession, ...]: + serving: Final = frozenset(session.pid for session in opened) + if len(opened) >= OPEN_SESSIONS and len(serving) == len(workers): + return opened + assert len(opened) < MAX_OPEN_SESSIONS, f"all {len(opened)} sessions landed on workers {sorted(serving)}" + return await _sessions_on_every_worker(url, key, workers, (*opened, await _open_session(url, key, workers))) + + +async def _sessions_through_worker_kill( + url: str, key: str, workers: tuple[psutil.Process, ...] +) -> tuple[tuple[Session, bool], ...]: + opened: Final = await _sessions_on_every_worker(url, key, workers, ()) + sockets: Final = tuple(session.socket for session in opened) + try: + on_killed_worker: Final = tuple(session.pid == workers[0].pid for session in opened) + workers[0].kill() + await asyncio.to_thread(eventually, lambda: _exited(workers[0]), bool, graceful_stop_seconds()) + turns: Final = (_turn_or_close(socket, label) for socket, label in zip(sockets, _labels(len(sockets)))) + sessions: Final = await asyncio.wait_for(asyncio.gather(*turns), 90) + return tuple(zip(sessions, on_killed_worker)) + finally: + await asyncio.gather(*(socket.close() for socket in sockets)) + + +def test_worker_kill_leaves_the_surviving_sessions_with_only_their_own_frames(gateway: Gateway, tmp_path: Path) -> None: + with gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _spoken()) + fresh: Final = _scripted(scenario, _spoken()) + key: Final = scenario.key() + fresh_key: Final = scenario.key() + with owned_proxy_process(gateway, tmp_path, {}, workers=WORKERS) as owned: + model: Final = _named_deployment(owned.gateway, scenario, handle.scenario_id) + fresh_model: Final = _named_deployment(owned.gateway, scenario, fresh.scenario_id) + root: Final = psutil.Process(owned.process.pid) + workers: Final = _workers(root) + assert len(workers) == WORKERS, [process.pid for process in workers] + outcomes: Final = asyncio.run( + _sessions_through_worker_kill(_session_url(_owned_url(owned), f"model={model}"), key, workers) + ) + served: Final = tuple(label for label, (_, killed) in zip(_labels(len(outcomes)), outcomes) if not killed) + assert [session.types for session, _ in outcomes] == [ + ("closed",) if killed else ("session.updated", TRANSCRIPT_COMPLETED, RESPONSE_DONE) + for _, killed in outcomes + ], outcomes + sent: Final = _sent(_observed(gateway.upstream_url, handle.scenario_id)) + assert _frame_counts(sent) == _frame_counts(_every_client_frame(served)), sent + with httpx.Client(base_url=_owned_url(owned), timeout=graceful_stop_seconds(), trust_env=False) as client: + readiness: Final = client.get("/health/readiness") + assert readiness.status_code == 200, readiness.text + eventually( + lambda: tuple(process.pid for process in _workers(root)), + lambda pids: len(pids) == WORKERS and workers[0].pid not in pids, + seconds=graceful_stop_seconds(), + ) + fresh_session: Final = _drive( + _owned_url(owned), f"model={fresh_model}", fresh_key, _client_turns(PUSH_TO_TALK, ("fresh",)) + ) + _assert_client_turns( + fresh_session, + _observed(gateway.upstream_url, fresh.scenario_id), + fresh.scenario_id, + PUSH_TO_TALK, + ("fresh",), + ) + _assert_billed(fresh_key, 1) + _assert_billed(key, len(served)) + + +def test_list_mode_transcript_guardrail_still_turns_auto_response_off_and_answers_the_turn( + guard_proxy: OwnedProxy, +) -> None: + with guard_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _spoken()) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id) + session: Final = _drive( + _owned_url(guard_proxy), + f"model={model}&guardrails={MIXED_GUARDRAIL}", + key, + (Step((_update_frame(VOICE_UPDATE), COMMIT), RESPONSE_DONE),), + ) + observed: Final = _observed(guard_proxy.gateway.upstream_url, handle.scenario_id) + assert session.types == ( + "session.created", + "session.updated", + "error", + TRANSCRIPT_COMPLETED, + RESPONSE_DONE, + ), session + assert session.transcript_fields == ({"transcript": TRANSCRIPT},), session + assert session.error_messages == (TURN_DETECTION_TYPE_MISSING,), session + assert _sent(observed) == ( + _update_frame(INJECTED_UPDATE), + _update_frame(VOICE_UPDATE_FORWARDED), + COMMIT, + PROXY_RESPONSE_CREATE, + ), _sent(observed) + _assert_billed(key, 1) diff --git a/tests/integration/providers/test_strands_decider_agentcore_wire.py b/tests/integration/providers/test_strands_decider_agentcore_wire.py new file mode 100644 index 00000000000..bf5da453137 --- /dev/null +++ b/tests/integration/providers/test_strands_decider_agentcore_wire.py @@ -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[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 diff --git a/tests/integration/routing/test_per_model_rate_limit_fallbacks.py b/tests/integration/routing/test_per_model_rate_limit_fallbacks.py new file mode 100644 index 00000000000..428a485f397 --- /dev/null +++ b/tests/integration/routing/test_per_model_rate_limit_fallbacks.py @@ -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"}) diff --git a/tests/integration/routing/test_per_model_rate_limit_fallbacks_lifecycle.py b/tests/integration/routing/test_per_model_rate_limit_fallbacks_lifecycle.py new file mode 100644 index 00000000000..4361843b2c9 --- /dev/null +++ b/tests/integration/routing/test_per_model_rate_limit_fallbacks_lifecycle.py @@ -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,)}) diff --git a/tests/integration/sdk/test_interception_converted_stream_options_sdk.py b/tests/integration/sdk/test_interception_converted_stream_options_sdk.py new file mode 100644 index 00000000000..17876fe605b --- /dev/null +++ b/tests/integration/sdk/test_interception_converted_stream_options_sdk.py @@ -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) diff --git a/tests/integration/sdk/test_openai_sdk_client_sessions_wire.py b/tests/integration/sdk/test_openai_sdk_client_sessions_wire.py new file mode 100644 index 00000000000..e8ec1a37082 --- /dev/null +++ b/tests/integration/sdk/test_openai_sdk_client_sessions_wire.py @@ -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 diff --git a/tests/integration/streaming/test_interception_converted_stream_options.py b/tests/integration/streaming/test_interception_converted_stream_options.py new file mode 100644 index 00000000000..46a4211b121 --- /dev/null +++ b/tests/integration/streaming/test_interception_converted_stream_options.py @@ -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 diff --git a/tests/integration/translation/decisions/bases/cloudflare.py b/tests/integration/translation/decisions/bases/cloudflare.py index ae86a2ccd87..3efc96eeb87 100644 --- a/tests/integration/translation/decisions/bases/cloudflare.py +++ b/tests/integration/translation/decisions/bases/cloudflare.py @@ -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", diff --git a/tests/integration/translation/decisions/bases/openrouter.py b/tests/integration/translation/decisions/bases/openrouter.py index b65acff6b06..ed39bc1a55b 100644 --- a/tests/integration/translation/decisions/bases/openrouter.py +++ b/tests/integration/translation/decisions/bases/openrouter.py @@ -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", diff --git a/tests/integration/translation/decisions/bases/perplexity.py b/tests/integration/translation/decisions/bases/perplexity.py index 02b1de51f24..9e29ceab8cc 100644 --- a/tests/integration/translation/decisions/bases/perplexity.py +++ b/tests/integration/translation/decisions/bases/perplexity.py @@ -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", diff --git a/tests/integration/translation/decisions/bases/strands_decider.py b/tests/integration/translation/decisions/bases/strands_decider.py index 8a1a51f8c3e..2950b090737 100644 --- a/tests/integration/translation/decisions/bases/strands_decider.py +++ b/tests/integration/translation/decisions/bases/strands_decider.py @@ -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", diff --git a/tests/integration/translation/decisions/bases/typesafe.py b/tests/integration/translation/decisions/bases/typesafe.py index 5da5b05cb09..7d47ce6b082 100644 --- a/tests/integration/translation/decisions/bases/typesafe.py +++ b/tests/integration/translation/decisions/bases/typesafe.py @@ -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", diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py b/tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py index 79fc0574b3b..21e0538897d 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_cloudflare.py @@ -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) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py b/tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py index a74cb23f6b4..973526d7272 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_openrouter.py @@ -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) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py b/tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py index e341d24ba53..d9662ec4b9f 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_perplexity.py @@ -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) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py b/tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py index cae7ca70c78..2d2cd9caaf5 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_strands_decider.py @@ -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) diff --git a/tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py b/tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py index 302636e6fb4..c1111cc6ba9 100644 --- a/tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py +++ b/tests/integration/translation/decisions/basic/test_decisions_basic_typesafe.py @@ -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) diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index dc66852d214..b3cd7bd01f9 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -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" diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index d3b04c8bb6f..9e3945f11b1 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -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 diff --git a/tests/test_litellm_rust/test_inference.py b/tests/test_litellm_rust/test_inference.py index 8211b50cb8d..83a01e8ad4d 100644 --- a/tests/test_litellm_rust/test_inference.py +++ b/tests/test_litellm_rust/test_inference.py @@ -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( diff --git a/tests/typing/openai_sdk_compat.py b/tests/typing/openai_sdk_compat.py new file mode 100644 index 00000000000..03f44db6a9f --- /dev/null +++ b/tests/typing/openai_sdk_compat.py @@ -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) diff --git a/tests/typing/pyrightconfig.json b/tests/typing/pyrightconfig.json new file mode 100644 index 00000000000..0910ca64f5b --- /dev/null +++ b/tests/typing/pyrightconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../pyrightconfig.json", + "include": ["openai_sdk_compat.py"], + "exclude": [], + "extraPaths": ["../.."], + "pythonVersion": "3.10", + "reportMissingImports": true +} diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 8f05a492119..ebe0eedea7f 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -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", diff --git a/tests/unit/integrations/code_interpreter_interception/test_handler.py b/tests/unit/integrations/code_interpreter_interception/test_handler.py index 9d308ac1989..c037b6be80d 100644 --- a/tests/unit/integrations/code_interpreter_interception/test_handler.py +++ b/tests/unit/integrations/code_interpreter_interception/test_handler.py @@ -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={ diff --git a/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py index 57f4557c6f7..ea0416cecbd 100644 --- a/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py @@ -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 diff --git a/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py index bb55ba57680..064b94942a3 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_interception_handler.py @@ -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. diff --git a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py index 0d7f735e6e5..e891706df13 100644 --- a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -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 [] + ) diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index 9f9f0d340d4..2a1c9d66965 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -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, ) diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_output_config_passthrough.py b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_output_config_passthrough.py index 282373d4b6d..03553e65ad4 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_output_config_passthrough.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_output_config_passthrough.py @@ -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" diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index 9daa60bbf88..eb2d62a8d71 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -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, diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 25d2d11b5ca..347840b5c5d 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -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() diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index 3774545488d..053ec7e141e 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -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, ) diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index fb605c6da7a..3c7dff34296 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -1,9 +1,10 @@ -from unittest.mock import MagicMock, call, patch +from typing import Final +from unittest.mock import MagicMock, patch import httpx import openai import pytest - +import respx import litellm from litellm.litellm_core_utils.token_counter import token_counter @@ -413,3 +414,124 @@ def test_is_openai_backed_api_base_decides_by_hostname_only(api_base, expected): assert is_openai_backed_api_base(api_base) is expected +def _sdk_api_client( + api: str, + is_async: bool, + timeout: float | httpx.Timeout | openai.Timeout | None, + client: openai.OpenAI | openai.AsyncOpenAI | None = None, +) -> openai.OpenAI | openai.AsyncOpenAI | None: + from litellm.llms.azure.common_utils import BaseAzureLLM + from litellm.llms.openai.fine_tuning.handler import OpenAIFineTuningAPI + from litellm.llms.openai.image_variations.handler import OpenAIImageVariationsHandler + from litellm.llms.openai.openai import OpenAIAssistantsAPI, OpenAIBatchesAPI, OpenAIFilesAPI + + kwargs: Final = { + "api_key": "transport-only", + "api_base": "https://sdk-default.example/v1", + "timeout": timeout, + "max_retries": 0, + "organization": None, + "client": client, + } + if api == "assistants": + assistant_factory: Final = OpenAIAssistantsAPI() + return ( + assistant_factory.async_get_openai_client(**kwargs) + if is_async + else assistant_factory.get_openai_client(**kwargs) + ) + if api == "image_variations": + variation_factory: Final = OpenAIImageVariationsHandler() + params: Final = { + "api_key": "transport-only", + "base_url": kwargs["api_base"], + "timeout": timeout, + } + return ( + variation_factory.get_async_client(client=client, init_client_params=params) + if is_async + else variation_factory.get_sync_client(client=client, init_client_params=params) + ) + if api == "azure_gateway": + return BaseAzureLLM()._init_azure_client_for_cloudflare_ai_gateway( + api_base="https://sdk-default.example", + model="deployment", + api_version="2024-02-01", + max_retries=0, + timeout=timeout, + litellm_params={}, + api_key="transport-only", + azure_ad_token=None, + azure_ad_token_provider=None, + acompletion=is_async, + client=client, + ) + factory: Final = {"files": OpenAIFilesAPI, "batches": OpenAIBatchesAPI, "fine_tuning": OpenAIFineTuningAPI}[api]() + return factory.get_openai_client(**kwargs, _is_async=is_async) + + +async def _list_model_ids(sdk_client: openai.OpenAI | openai.AsyncOpenAI) -> list[str]: + page: Final = ( + await sdk_client.models.list() if isinstance(sdk_client, openai.AsyncOpenAI) else sdk_client.models.list() + ) + return [model.id for model in page.data] + + +@pytest.mark.parametrize("api", ["files", "batches", "assistants", "fine_tuning", "image_variations", "azure_gateway"]) +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_sdk_api_factories_keep_httpx_transport_and_request_timeouts(api: str, is_async: bool) -> None: + timeout: Final = openai.Timeout(connect=1, read=7, write=2, pool=3) + with respx.mock(assert_all_called=True) as mock: + mock.get(host="sdk-default.example", path__regex=r".*/models$").mock( + return_value=httpx.Response(307, headers={"location": "https://sdk-default.example/v1/moved-models"}) + ) + moved: Final = mock.get(host="sdk-default.example", path="/v1/moved-models").mock( + return_value=httpx.Response( + 200, + json={"object": "list", "data": [{"id": "sdk-compat", "object": "model", "created": 0, "owned_by": "t"}]}, + ) + ) + sdk_client: Final = _sdk_api_client(api, is_async, timeout) + assert sdk_client is not None + try: + assert _sdk_api_client(api, is_async, 19, client=sdk_client) is sdk_client + model_ids: Final = await _list_model_ids(sdk_client) + finally: + if is_async: + await sdk_client.close() + else: + sdk_client.close() + assert model_ids == ["sdk-compat"] + assert moved.calls.last.request.extensions["timeout"] == {"connect": 1.0, "read": 7.0, "write": 2.0, "pool": 3.0} + + +@pytest.mark.parametrize("backend", ["httpx", "sdk_default"]) +@pytest.mark.asyncio +async def test_azure_gateway_and_image_variations_use_the_callers_async_session( + backend: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from importlib import import_module + from io import BytesIO + + from openai._types import Response as SDKResponse + + from litellm.images.main import aimage_variation + + http_module: Final = httpx if backend == "httpx" else import_module(SDKResponse.__module__) + transport: Final = http_module.MockTransport( + lambda request: http_module.Response( + 200, json={"created": 0, "data": [{"url": "https://example.com/image.png"}]} + ) + ) + async with http_module.AsyncClient(transport=transport) as session: + monkeypatch.setattr(litellm, "aclient_session", session) + monkeypatch.setattr(litellm, "client_session", None) + azure_client: Final = _sdk_api_client("azure_gateway", True, 7) + assert azure_client is not None + assert azure_client._client is session + response: Final = await aimage_variation( + image=BytesIO(b"image-bytes"), api_key="transport-only", api_base="https://sdk-default.example/v1" + ) + assert response.data[0].url == "https://example.com/image.png" + assert not session.is_closed diff --git a/tests/unit/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py index 7d78a2c5238..d0cf52d0e28 100644 --- a/tests/unit/llms/openai/test_openai_workload_identity.py +++ b/tests/unit/llms/openai/test_openai_workload_identity.py @@ -1,5 +1,6 @@ import json import sys +from functools import partial from pathlib import Path from typing import Final @@ -36,6 +37,24 @@ CHAT_COMPLETION_BODY: Final = { } +@pytest.fixture(autouse=True) +def mock_sdk_token_exchange_transport(monkeypatch: pytest.MonkeyPatch) -> None: + native_client_factory: Final = getattr(sys.modules.get("openai.auth._workload"), "DefaultHttpx2Client", None) + if native_client_factory is not None: + import httpx2 + + def handle_exchange(request: httpx2.Request) -> httpx2.Response: + response: Final = respx.mock.handler( + httpx.Request(request.method, str(request.url), headers=dict(request.headers), content=request.content) + ) + return httpx2.Response(response.status_code, headers=dict(response.headers), content=response.content) + + monkeypatch.setattr( + "openai.auth._workload.DefaultHttpx2Client", + partial(native_client_factory, transport=httpx2.MockTransport(handle_exchange), trust_env=False), + ) + + @pytest.fixture def wif_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> OpenAIWorkloadIdentityConfig: token_file: Final = tmp_path / "subject_token.jwt" @@ -602,7 +621,6 @@ class TestDiscoverModels: assert models_route.calls.last.request.headers["Authorization"] == "Bearer None" assert not exchange_route.called - @respx.mock def test_empty_static_key_never_borrows_the_env_key(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("OPENAI_API_KEY", "sk-env-key-that-must-stay-home") diff --git a/tests/unit/llms/strands_decider/__init__.py b/tests/unit/llms/strands_decider/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/strands_decider/decisions/__init__.py b/tests/unit/llms/strands_decider/decisions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/strands_decider/decisions/test_strands_decider_decisions_transformation.py b/tests/unit/llms/strands_decider/decisions/test_strands_decider_decisions_transformation.py new file mode 100644 index 00000000000..4d249e59233 --- /dev/null +++ b/tests/unit/llms/strands_decider/decisions/test_strands_decider_decisions_transformation.py @@ -0,0 +1,256 @@ +import json +from collections.abc import Mapping +from typing import Final + +import httpx +import pytest +from botocore.credentials import Credentials +from pydantic import TypeAdapter + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.decisions.transformation import systemone_request_to_ir +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.strands_decider.decisions.transformation import ( + AGENTCORE_SESSION_HEADER, + StrandsDeciderDecisionsConfig, +) +from litellm.types.decisions import DecisionsIRResponse, DecisionsRequestBody +from litellm.types.llms.bedrock import AwsAuthParams + +_ARN: Final = "arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/qa_decider-AbC123xyz0" +_BODY: Final[Mapping[str, object]] = { + "model": "strands-decider-2B-hobson-v21", + "state": "My order arrived broken.", + "questions": {"is_complaint": {"type": "noul", "instructions": "Is this a complaint?"}}, +} +_IR_REQUEST: Final = systemone_request_to_ir( + TypeAdapter(DecisionsRequestBody).validate_python({"state": _BODY["state"], "questions": _BODY["questions"]}) +) + + +class _RecordingAWS(BaseAWSLLM): + def __init__(self) -> None: + super().__init__() + self.resolved: tuple[tuple[AwsAuthParams, str | None], ...] = () + + def resolve_credentials(self, auth_params: AwsAuthParams, aws_region_name: str | None) -> Credentials: + self.resolved = (*self.resolved, (auth_params, aws_region_name)) + return Credentials(access_key="AKIDEXAMPLE", secret_key="example-secret") + + +def _sign( + config: StrandsDeciderDecisionsConfig, + api_base: str, + headers: Mapping[str, str], + api_key: str | None = None, + litellm_params: Mapping[str, object] | None = None, +) -> tuple[Mapping[str, str], bytes | None]: + return config.sign_request( + headers=headers, + url=config.get_complete_url(api_base=api_base, model="strands-decider-2B-hobson-v21"), + api_base=api_base, + body=_BODY, + api_key=api_key, + litellm_params=litellm_params or {}, + ) + + +@pytest.mark.parametrize( + ("arn", "expected_url"), + [ + ( + _ARN, + "https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/" + "arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A123456789012%3Aruntime%2Fqa_decider-AbC123xyz0/invocations", + ), + ( + "arn:aws-cn:bedrock-agentcore:cn-north-1:123456789012:runtime/qa_decider-AbC123xyz0", + "https://bedrock-agentcore.cn-north-1.amazonaws.com.cn/runtimes/" + "arn%3Aaws-cn%3Abedrock-agentcore%3Acn-north-1%3A123456789012%3Aruntime%2Fqa_decider-AbC123xyz0" + "/invocations", + ), + ( + "arn:aws-us-gov:bedrock-agentcore:us-gov-west-1:123456789012:runtime/qa_decider-AbC123xyz0", + "https://bedrock-agentcore.us-gov-west-1.amazonaws.com/runtimes/" + "arn%3Aaws-us-gov%3Abedrock-agentcore%3Aus-gov-west-1%3A123456789012%3Aruntime%2Fqa_decider-AbC123xyz0" + "/invocations", + ), + ], +) +def test_runtime_arn_targets_the_regional_invoke_agent_runtime_url(arn: str, expected_url: str) -> None: + assert StrandsDeciderDecisionsConfig().get_complete_url(api_base=arn, model="any") == expected_url + + +@pytest.mark.parametrize( + "api_base", + [ + "https://strands.example", + "arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/qa_decider-AbC123xyz0/runtime-endpoint/DEFAULT", + "arn:aws:bedrock-agentcore:evil.example/x:123456789012:runtime/qa_decider-AbC123xyz0", + "arn:aws:bedrock:us-west-2:123456789012:runtime/qa_decider-AbC123xyz0", + ], +) +def test_other_api_bases_keep_the_plain_systemone_request(api_base: str) -> None: + aws: Final = _RecordingAWS() + config: Final = StrandsDeciderDecisionsConfig(aws=aws) + headers: Final = {"Content-Type": "application/json"} + + assert config.get_complete_url(api_base=api_base, model="any") == f"{api_base}/v1/systemone" + assert _sign(config, api_base, headers) == (headers, None) + assert aws.resolved == () + + +def test_sign_request_signs_the_body_for_the_arn_region_with_deployment_credentials() -> None: + aws: Final = _RecordingAWS() + config: Final = StrandsDeciderDecisionsConfig(aws=aws) + + signed_headers, signed_body = _sign( + config, + _ARN, + {"Content-Type": "application/json"}, + litellm_params={"aws_profile_name": "qa", "aws_role_name": "decider-role", "aws_region_name": "us-east-1"}, + ) + + assert signed_body == json.dumps(_BODY).encode() + assert len(aws.resolved) == 1 + auth_params, region = aws.resolved[0] + assert (auth_params.aws_profile_name, auth_params.aws_role_name, region) == ("qa", "decider-role", "us-west-2") + authorization: Final = signed_headers["Authorization"] + assert authorization.startswith("AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/") + assert "/us-west-2/bedrock-agentcore/aws4_request" in authorization + assert AGENTCORE_SESSION_HEADER.lower() in authorization + + +def test_default_session_id_is_stable_per_runtime_and_within_the_agentcore_length_bounds() -> None: + other_arn: Final = _ARN.replace("AbC123xyz0", "ZyX987cba1") + + first, _ = _sign(StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), _ARN, {}) + again, _ = _sign(StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), _ARN, {}) + other, _ = _sign(StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), other_arn, {}) + + session_id: Final = first[AGENTCORE_SESSION_HEADER] + assert session_id == again[AGENTCORE_SESSION_HEADER] + assert session_id != other[AGENTCORE_SESSION_HEADER] + assert 33 <= len(session_id) <= 256 + + +_CALLER_RESERVED_HEADERS: Final = { + AGENTCORE_SESSION_HEADER.lower(): "caller-session-0123456789abcdef0123456789", + "X-Amzn-Bedrock-AgentCore-Runtime-User-Id": "someone-else", + "X-Amz-Security-Token": "caller-token", + "Host": "evil.example", +} + + +@pytest.mark.parametrize("api_key", [None, "runtime-jwt"]) +def test_caller_reserved_agentcore_headers_never_reach_the_runtime(api_key: str | None) -> None: + signed_headers, _ = _sign( + StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), + _ARN, + {**_CALLER_RESERVED_HEADERS, "X-Request-Source": "batch"}, + api_key=api_key, + ) + default_headers, _ = _sign(StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), _ARN, {}, api_key=api_key) + + forwarded: Final = {name.lower(): value for name, value in signed_headers.items()} + assert forwarded[AGENTCORE_SESSION_HEADER.lower()] == default_headers[AGENTCORE_SESSION_HEADER] + assert not any(value in forwarded.values() for value in _CALLER_RESERVED_HEADERS.values()) + assert "x-amzn-bedrock-agentcore-runtime-user-id" not in forwarded + assert "host" not in forwarded + assert forwarded["x-request-source"] == "batch" + + +def test_deployment_session_param_picks_the_session_over_the_default_and_any_caller_header() -> None: + replica_session: Final = "strands-decider-replica-a-0000000000" + + signed_headers, _ = _sign( + StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), + _ARN, + _CALLER_RESERVED_HEADERS, + litellm_params={"agentcore_runtime_session_id": replica_session}, + ) + + session_headers: Final = [ + value for name, value in signed_headers.items() if name.lower() == "x-amzn-bedrock-agentcore-runtime-session-id" + ] + assert session_headers == [replica_session] + assert AGENTCORE_SESSION_HEADER.lower() in signed_headers["Authorization"] + + +def test_api_key_on_a_runtime_arn_sends_the_bearer_token_without_sigv4() -> None: + aws: Final = _RecordingAWS() + + signed_headers, signed_body = _sign( + StrandsDeciderDecisionsConfig(aws=aws), _ARN, {"Authorization": "Bearer runtime-jwt"}, api_key="runtime-jwt" + ) + + assert signed_body is None + assert signed_headers["Authorization"] == "Bearer runtime-jwt" + assert AGENTCORE_SESSION_HEADER in signed_headers + assert aws.resolved == () + + +def _response(url: str, payload: object) -> httpx.Response: + return httpx.Response(200, json=payload, request=httpx.Request("POST", url)) + + +def _transform(api_base: str, payload: object) -> DecisionsIRResponse: + config: Final = StrandsDeciderDecisionsConfig() + url: Final = config.get_complete_url(api_base=api_base, model="strands-decider-2B-hobson-v21") + return config.transform_decisions_response( + model="strands-decider-2B-hobson-v21", + custom_llm_provider="strands_decider", + raw_response=_response(url, payload), + request=_IR_REQUEST, + ) + + +@pytest.mark.parametrize( + ("code", "status_code"), + [("bad_request", 400), ("invalid_request", 400), ("loading", 503), ("inference_error", 500), ("new_code", 500)], +) +def test_runtime_error_envelope_raises_with_the_mapped_status(code: str, status_code: int) -> None: + with pytest.raises(BaseLLMException) as raised: + _transform(_ARN, {"error": {"code": code, "message": "at most 16 questions"}}) + + assert raised.value.status_code == status_code + assert code in raised.value.message + assert "at most 16 questions" in raised.value.message + + +@pytest.mark.parametrize("code", ["bad_request", "loading", "inference_error"]) +def test_plain_api_base_error_envelope_keeps_the_unexpected_response_error(code: str) -> None: + with pytest.raises(BaseLLMException) as raised: + _transform("https://strands.example", {"error": {"code": code, "message": "at most 16 questions"}}) + + assert raised.value.status_code == 500 + assert raised.value.message.startswith("Decisions provider 'strands_decider' returned an unexpected response") + + +@pytest.mark.parametrize( + "url", + [ + "https://strands.example/runtimes/arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A123456789012%3Aruntime" + "%2Fqa_decider-AbC123xyz0/invocations", + "https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/not-an-arn/invocations", + ], +) +def test_error_envelope_from_a_url_that_is_not_the_runtime_keeps_the_unexpected_response_error(url: str) -> None: + with pytest.raises(BaseLLMException) as raised: + StrandsDeciderDecisionsConfig().transform_decisions_response( + model="strands-decider-2B-hobson-v21", + custom_llm_provider="strands_decider", + raw_response=_response(url, {"error": {"code": "bad_request", "message": "at most 16 questions"}}), + request=_IR_REQUEST, + ) + + assert raised.value.status_code == 500 + + +@pytest.mark.parametrize("payload", [["not", "an", "object"], {"error": "unstructured"}]) +def test_runtime_reply_that_is_not_an_error_envelope_keeps_the_unexpected_response_error(payload: object) -> None: + with pytest.raises(BaseLLMException) as raised: + _transform(_ARN, payload) + + assert raised.value.status_code == 500 + assert raised.value.message.startswith("Decisions provider 'strands_decider' returned an unexpected response") diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index ad94defebd7..adae7872430 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -26,6 +26,24 @@ from litellm.proxy._types import ( from litellm.types.agents import AgentCaller +@pytest.mark.parametrize("headers,expected", [ + ({"Authorization": "Bearer caller.jwt.signature"}, "caller.jwt.signature"), + ({"AUTHORIZATION": "bEaReR caller.jwt.signature"}, "caller.jwt.signature"), + ({"authorization": "Bearer "}, ""), + ({"authorization": "Basic credentials"}, None), + ({"authorization-extra": "Bearer wrong.jwt.signature"}, None), + ({"x-litellm-api-key": "sk-key"}, None), + ({}, None), +]) +def test_incoming_guardrail_bearer_requires_exact_authorization_header( + headers: dict[str, str], expected: str | None, +) -> None: + token: Final = MCPRequestHandler.get_incoming_bearer_token(headers) + assert (token.get_secret_value() if token is not None else None) == expected + if expected: + assert expected not in repr(token) + + @pytest.mark.asyncio class TestMCPRequestHandler: @pytest.mark.parametrize( diff --git a/tests/unit/proxy/_experimental/mcp_server/test_contracts.py b/tests/unit/proxy/_experimental/mcp_server/test_contracts.py index e13ecdfcce9..db877bd4395 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_contracts.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_contracts.py @@ -1,11 +1,24 @@ from dataclasses import FrozenInstanceError +from typing import Final import pytest +from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.operations import prepare_context from litellm.proxy._types import UserAPIKeyAuth +def test_guardrail_bearer_stays_out_of_egress_and_context_representations() -> None: + token: Final = SecretStr("caller.jwt.signature") + context: Final = prepare_context(incoming_bearer_token=token) + assert context.incoming_bearer_token is token + assert token.get_secret_value() not in repr(context) + assert token not in context.legacy_auth() + assert context.raw_headers is None + assert context.oauth2_headers is None + assert context.mcp_server_auth_headers is None + + def test_operation_context_isolates_nested_headers_and_caller_permissions(): caller = UserAPIKeyAuth(user_id="alpha", models=["allowed"]) caller.mcp_admitted_user_subject = True diff --git a/tests/unit/proxy/_experimental/mcp_server/test_operations.py b/tests/unit/proxy/_experimental/mcp_server/test_operations.py index 25c80398e46..32a7c3fbe58 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_operations.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_operations.py @@ -1,9 +1,10 @@ import asyncio -from collections.abc import Sequence +from collections.abc import Callable, Sequence from typing import Final from unittest.mock import AsyncMock, patch import pytest +from mcp.server.context import ServerRequestContext from mcp.types import GetPromptRequest, GetPromptRequestParams, GetPromptResult from mcp.types import Tool as MCPTool @@ -43,6 +44,116 @@ async def _served_catalog_tool() -> str: return "ok" +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["direct", "virtual", "proxy"]) +@pytest.mark.parametrize("verdict", ["allow", "block", "obo_error", "evaluate_error", "missing"]) +async def test_mapped_caller_bearer_reaches_agent_evaluation( + monkeypatch: pytest.MonkeyPatch, _mcp_request_ctx: Callable[..., ServerRequestContext], + mode: str, verdict: str, caplog: pytest.LogCaptureFixture, +) -> None: + from urllib.parse import parse_qs + + import httpx + from mcp.server.auth.middleware.auth_context import auth_context_var + from mcp.types import CallToolRequest, CallToolRequestParams, CallToolResult, TextContent + from starlette.requests import Request + + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.proxy._experimental.mcp_server import server as gateway + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_proxy_mode + from litellm.proxy._experimental.mcp_server.tool_search import mcp_proxy_tool_id, with_mcp_proxy_identity + from litellm.proxy.guardrails.guardrail_hooks.agent_365.agent_365 import Agent365Guardrail + + bearer: Final = "signed.caller.assertion" + request: Final = Request({ + "type": "http", "method": "POST", "path": "/mcp", + "headers": [] if verdict == "missing" else [(b"authorization", f"Bearer {bearer}".encode())], + }) + auth: Final = UserAPIKeyAuth( + api_key="stored-key-hash", user_id="mapped-caller", + object_permission={"object_permission_id": "mapped", "mcp_servers": ["guarded"], "mcp_tool_search_enabled": mode == "virtual"}, + ) + oauth, raw, _, _ = MCPRequestHandler.scrub_gateway_admission_credentials( + admitted=False, admitted_credential=bearer, + oauth2_headers={"Authorization": f"Bearer {bearer}"}, + raw_headers=dict(request.headers), mcp_auth_header=None, mcp_server_auth_headers=None, + ) + assert "authorization" not in raw + provider_requests: Final[asyncio.Queue[str]] = asyncio.Queue() + + def provider(exchange: httpx.Request) -> httpx.Response: + provider_requests.put_nowait(exchange.url.path) + if exchange.url.path.endswith("/token"): + assert parse_qs(exchange.content.decode())["assertion"] == [bearer] + if verdict == "obo_error": + return httpx.Response(400, json={"error": "invalid_grant"}) + return httpx.Response(200, json={"access_token": "delegated-token", "expires_in": 3600}) + assert exchange.headers["authorization"] == "Bearer delegated-token" + if verdict == "evaluate_error": + return httpx.Response(503, json={"error": "unavailable"}) + return httpx.Response(200, json={ + "allowed": verdict == "allow", "defender": {"status": "Evaluated", "verdict": "Allow" if verdict == "allow" else "Block"}, + }) + + executions: Final[asyncio.Queue[str]] = asyncio.Queue() + + async def tool() -> str: + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import request_extra_headers + + assert bearer not in str(request_extra_headers.get()) + executions.put_nowait("executed") + return "ok" + + handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(provider)) + guardrail: Final = Agent365Guardrail( + guardrail_name="agent-evaluation", tenant_id="tenant", client_id="gateway", client_secret="secret", + async_handler=handler, event_hook="pre_mcp_call", default_on=True, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + managed: Final = MCPServer( + server_id="guarded", name="guarded", transport=MCPTransport.http, + spec_path="/catalog.yaml", allow_all_keys=True, + ) + monkeypatch.setitem(operations.global_mcp_server_manager.registry, managed.server_id, managed) + global_mcp_tool_registry.register_tool( + name="guarded-echo", description="Echo", input_schema={"type": "object"}, handler=tool, + ) + previous: Final = auth_context_var.set(None) + proxy_mode: Final = _mcp_proxy_mode.set(mode == "proxy") + tool_id: Final = mcp_proxy_tool_id(with_mcp_proxy_identity( + MCPTool(name="guarded-echo", input_schema={"type": "object"}), managed.server_id, + )) + params: Final = { + "direct": CallToolRequestParams(name="guarded-echo", arguments={}), + "virtual": CallToolRequestParams(name="mcp_tool_call", arguments={"tool_name": "guarded-echo", "arguments": {}}), + "proxy": CallToolRequestParams(name="call_tool", arguments={"tool_id": tool_id, "arguments": {}}), + }[mode] + try: + gateway.set_auth_context( + auth, mcp_servers=[managed.server_id], oauth2_headers=oauth, + raw_headers={"authorization": "Bearer previous.caller.assertion"} if verdict == "missing" else raw, + ) + async with gateway._legacy_operation_context(_mcp_request_ctx(request=request), trace=False) as context: + result: Final = await GatewayOperations().execute( + CallToolRequest(params=params), context, + ) + assert isinstance(result, CallToolResult) + assert result.is_error is (verdict != "allow") + assert executions.qsize() == (1 if verdict == "allow" else 0) + assert provider_requests.qsize() == (0 if verdict == "missing" else 1 if verdict == "obo_error" else 2) + assert bearer not in caplog.text + assert bearer not in repr(context) + if verdict == "allow": + assert isinstance(result.content[0], TextContent) + assert result.content[0].text == "ok" + finally: + auth_context_var.reset(previous) + _mcp_proxy_mode.reset(proxy_mode) + global_mcp_tool_registry.unregister_tools_with_prefix("guarded-") + await handler.close() + + @pytest.mark.asyncio @pytest.mark.parametrize("surface", ["mcp", "rest"]) @pytest.mark.parametrize("restriction", ["key", "server"]) diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 8be1cec0a0c..d05af0b388b 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -3025,6 +3025,22 @@ class TestIsRequestBodySafeBlocksBedrockProjectOverride: ) +def test_agentcore_runtime_session_id_in_request_body_is_rejected_without_admin_opt_in(): + body: Final = {"model": "strands-decider", "agentcore_runtime_session_id": "caller-session-0123456789abcdef0123"} + + with pytest.raises(ValueError, match="agentcore_runtime_session_id"): + is_request_body_safe(request_body=body, general_settings={}, llm_router=None, model="strands-decider") + assert ( + is_request_body_safe( + request_body=body, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="strands-decider", + ) + is True + ) + + class TestIsRequestBodySafeBlocksClaudePlatformWorkspaceOverride: @pytest.mark.parametrize( "alias", ["workspace_id", "aws_workspace_id", "anthropic_workspace_id", "anthropic-workspace-id"] diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index 6ce15636f1e..ed78fc7063a 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -575,6 +575,68 @@ def test_openai_format_decisions_reach_an_openai_deployment_unchanged_including_ assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected_cost) +@pytest.mark.parametrize("safety_identifier", (7, ["end-user-1"]), ids=("numeric", "list")) +@pytest.mark.parametrize("deployment_drops_params", (False, True), ids=("strict", "drop_params")) +def test_a_non_string_safety_identifier_is_refused_unless_the_deployment_drops_params( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + safety_identifier: object, + deployment_drops_params: bool, +) -> 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) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "decider", + "litellm_params": { + "model": "openai/gpt-6-luna", + "api_key": "k", + "drop_params": deployment_drops_params, + }, + } + ] + ), + ) + upstream: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond( + json={ + "model": "gpt-6-luna", + "answers": _OPENAI_FORMAT_ANSWERS[:1], + "usage": { + "input_tokens": _INPUT_TOKENS, + "output_tokens": _OUTPUT_TOKENS, + "total_tokens": _INPUT_TOKENS + _OUTPUT_TOKENS, + }, + } + ) + request_body: Final = { + "model": "decider", + "input": "The package arrived with a broken screen.", + "questions": _OPENAI_FORMAT_REQUEST["questions"][:1], + "safety_identifier": safety_identifier, + } + + response: Final = client.post("/v1/decisions", json=request_body) + + if not deployment_drops_params: + assert response.status_code == 400, response.text + assert "safety_identifier" in response.json()["error"]["message"] + assert not upstream.called + return + assert response.status_code == 200, response.text + assert json.loads(upstream.calls[0].request.content) == { + "model": "gpt-6-luna", + "input": "The package arrived with a broken screen.", + "questions": _OPENAI_FORMAT_REQUEST["questions"][:1], + } + + def _decisions_feature() -> LazyFeature: return next(feature for feature in LAZY_FEATURES if feature.name == "decisions") diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py index f3456f6b60c..6f4b1816150 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -23,7 +23,7 @@ Tests cover: import json import time -from typing import Optional +from typing import Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -41,7 +41,7 @@ from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( from litellm.proxy.spend_tracking.compression_savings import ( extract_compression_saved_tokens, ) -from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY +from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY, HEADROOM_STREAM_OPTIONS_KEY from litellm.types.utils import ( CallTypes, GenericGuardrailAPIInputs, @@ -2171,6 +2171,28 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple assert kwargs["stream"] is True +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", (CallTypes.acompletion, CallTypes.aresponses)) +async def test_pre_call_deployment_hook_moves_stream_options_off_the_forced_non_stream_call( + guardrail: HeadroomGuardrail, call_type: CallTypes +): + stream_options: Final = {"include_usage": True} + kwargs: Final[dict[str, object]] = { + "model": "gpt-4o", + "stream": True, + "stream_options": stream_options, + "tools": [_retrieve_tool_definition()], + } + + result: Final = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=call_type) + + assert result is not None + assert result["stream"] is False + assert "stream_options" not in result + assert result[HEADROOM_STREAM_OPTIONS_KEY] == stream_options + assert kwargs["stream_options"] == stream_options + + @pytest.mark.asyncio async def test_pre_call_deployment_hook_leaves_background_streams_alone(guardrail: HeadroomGuardrail): kwargs = { diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py index 954b9b99622..a6e84a9db0a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py @@ -1,15 +1,17 @@ +import asyncio import os import threading import time import uuid -from typing import List, cast +from typing import Final, List, Literal, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import HTTPException -from httpx import Request, Response import requests - +import responses +import respx +from fastapi import HTTPException +from httpx import HTTPStatusError, Request, Response import litellm from litellm import ModelResponse @@ -1153,3 +1155,130 @@ def test_get_jwt_gives_up_at_the_timeout_instead_of_blocking_the_event_loop(hang guardrail.refresh_jwt_func() assert [call.kwargs["timeout"] for call in get_jwt.call_args_list] == [2, 2] + + +@pytest.fixture(params=("v1", "v2")) +def hiddenlayer_jwt_guardrail( + request: pytest.FixtureRequest, httpx_transport: None +) -> HiddenlayerGuardrail | HiddenlayerGuardrailV2: + guardrail_cls: Final = HiddenlayerGuardrail if request.param == "v1" else HiddenlayerGuardrailV2 + with responses.RequestsMock() as auth: + auth.post( + "https://auth.example.test/oauth2/token?grant_type=client_credentials", json={"access_token": "expired"} + ) + return guardrail_cls( + guardrail_name="hiddenlayer", + api_id="id", + api_key="secret", + api_base="https://api.hiddenlayer.ai", + auth_url="https://auth.example.test", + timeout=2, + ) + + +@pytest.mark.parametrize("input_type", ("request", "response")) +async def test_jwt_refresh_allows_other_coroutines_to_run_before_it_finishes( + hiddenlayer_jwt_guardrail: HiddenlayerGuardrail | HiddenlayerGuardrailV2, + input_type: Literal["request", "response"], + respx_mock: respx.MockRouter, +) -> None: + guardrail: Final = hiddenlayer_jwt_guardrail + loop: Final = asyncio.get_running_loop() + refresh_started: Final = asyncio.Event() + release_refresh: Final = threading.Event() + successful_body: Final = ( + {"evaluation": {"action": "Allow"}} + if isinstance(guardrail, HiddenlayerGuardrail) + else {"messages": [{"role": "user", "content": "hi"}]} + if input_type == "request" + else {"choices": [{"message": {"role": "assistant", "content": "hi"}}]} + ) + + def refresh_token(request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + assert request.headers["Authorization"] == "Basic aWQ6c2VjcmV0" + loop.call_soon_threadsafe(refresh_started.set) + assert release_refresh.wait(timeout=5), "event loop could not run while JWT refresh was pending" + return 200, {}, '{"access_token": "fresh"}' + + def evaluate(request: Request) -> Response: + return Response( + 401 if request.headers["Authorization"] == "Bearer expired" else 200, + json=successful_body, + ) + + async def other_request() -> None: + await refresh_started.wait() + release_refresh.set() + + route: Final = respx_mock.route(method="POST", host="api.hiddenlayer.ai").mock(side_effect=evaluate) + peer: Final = asyncio.create_task(other_request()) + with responses.RequestsMock() as auth: + auth.add_callback( + responses.POST, + "https://auth.example.test/oauth2/token?grant_type=client_credentials", + callback=refresh_token, + content_type="application/json", + ) + try: + result, _ = await asyncio.wait_for( + asyncio.gather( + guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type=input_type), + peer, + ), + timeout=10, + ) + finally: + release_refresh.set() + peer.cancel() + await asyncio.gather(peer, return_exceptions=True) + assert len(auth.calls) == 1 + + assert result["texts"] == ["hi"] + assert guardrail.jwt_token == "fresh" + assert route.call_count == 2 + assert route.calls[0].request.headers["Authorization"] == "Bearer expired" + assert route.calls[1].request.headers["Authorization"] == "Bearer fresh" + + +@pytest.mark.parametrize("input_type", ("request", "response")) +async def test_jwt_refresh_exception_propagates_without_replacing_the_token( + hiddenlayer_jwt_guardrail: HiddenlayerGuardrail | HiddenlayerGuardrailV2, + input_type: Literal["request", "response"], + respx_mock: respx.MockRouter, +) -> None: + guardrail: Final = hiddenlayer_jwt_guardrail + refresh_error: Final = requests.exceptions.Timeout("token endpoint timed out") + route: Final = respx_mock.route(method="POST", host="api.hiddenlayer.ai").respond(401) + with responses.RequestsMock() as auth: + auth.post("https://auth.example.test/oauth2/token?grant_type=client_credentials", body=refresh_error) + with pytest.raises(requests.exceptions.Timeout, match="token endpoint timed out") as caught: + await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type=input_type) + + assert caught.value is refresh_error + assert guardrail.jwt_token == "expired" + assert route.call_count == 1 + assert route.calls[0].request.headers["Authorization"] == "Bearer expired" + + +@pytest.mark.parametrize("input_type", ("request", "response")) +async def test_second_detection_401_propagates_without_another_jwt_refresh( + hiddenlayer_jwt_guardrail: HiddenlayerGuardrail | HiddenlayerGuardrailV2, + input_type: Literal["request", "response"], + respx_mock: respx.MockRouter, +) -> None: + guardrail: Final = hiddenlayer_jwt_guardrail + route: Final = respx_mock.route(method="POST", host="api.hiddenlayer.ai").respond(401) + with responses.RequestsMock() as auth: + auth.post( + "https://auth.example.test/oauth2/token?grant_type=client_credentials", json={"access_token": "fresh"} + ) + with pytest.raises(HTTPStatusError) as caught: + await guardrail.apply_guardrail(inputs={"texts": ["hi"]}, request_data={}, input_type=input_type) + assert len(auth.calls) == 1 + + assert caught.value.response.status_code == 401 + assert caught.value.request.headers["Authorization"] == "Bearer fresh" + assert guardrail.jwt_token == "fresh" + assert route.call_count == 2 + assert route.calls[0].request.headers["Authorization"] == "Bearer expired" + assert route.calls[1].request.headers["Authorization"] == "Bearer fresh" diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index a2e43b3bc9c..89cb6bd1ef2 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -3,7 +3,7 @@ Unit Tests for the max parallel request limiter v1 for the proxy """ import itertools -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from datetime import datetime from typing import Final @@ -40,6 +40,39 @@ def _clock_rolling_over_after_first_read() -> Callable[[], datetime]: ) +@pytest.mark.parametrize( + "current,rpm_limit", + [ + (None, 0), + ({"current_requests": 0, "current_tpm": 0, "current_rpm": 1}, 1), + ], +) +@pytest.mark.asyncio +async def test_model_per_key_rate_limit_error_carries_descriptor_key( + current: Mapping[str, int] | None, rpm_limit: int +): + handler: Final = PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.check_key_in_limits( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data={"model": "gpt-4o-mini"}, + call_type="completion", + max_parallel_requests=10, + tpm_limit=100, + rpm_limit=rpm_limit, + current=dict(current) if current is not None else None, + request_count_api_key="test-key:model_per_key", + rate_limit_type="model_per_key", + values_to_update_in_cache=[], + ) + + assert exc_info.value.descriptor_key == "model_per_key" + + @pytest.mark.asyncio async def test_pre_call_hook_counts_a_cli_session_under_the_per_user_alias_not_the_login_token(): handler = PROXY_MaxParallelRequestsHandler( diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index d9cf8536912..55163c6240b 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7467,6 +7467,36 @@ def test_rate_limit_error_reports_reset_time_in_utc_on_a_non_utc_proxy(process_t ) +def test_per_model_rate_limit_error_carries_descriptor_key() -> None: + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + over_limit: Final[RateLimitResponse] = { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "model_per_key", + "descriptor_value": "gpt-4o-mini", + "limit_remaining": 0, + "rate_limit_type": "requests", + "current_limit": 2, + } + ], + } + + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._handle_rate_limit_error( + response=over_limit, + descriptors=[ + {"key": "model_per_key", "value": "gpt-4o-mini", "rate_limit": None} + ], + requested_model="gpt-4o-mini", + ) + + assert exc_info.value.descriptor_key == "model_per_key" + + def _resolve_alias_to_target(model: str) -> str | None: return "target" if model == "alias" else None diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index d2065d34bd3..b040cc861bf 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -2,6 +2,7 @@ import asyncio import copy import datetime import json +from collections.abc import Mapping from types import MappingProxyType, SimpleNamespace from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence from urllib.parse import unquote_plus @@ -7095,12 +7096,16 @@ class TestPreCallWithFallbacksOnLocalRateLimit: user_api_key_dict: ProxyUserAPIKeyAuth, fallbacks: list[dict[str, list[str]]], model_guardrails: dict[str, list[str]] | None = None, + untagged_rate_limited_models: frozenset[str] = frozenset(), ) -> tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]]: """Real v3 limiter (the default ``parallel_request_limiter``) wired in through the ``proxy_logging_obj`` seam, so ``common_processing_pre_call_logic`` runs for real: - ``add_litellm_data_to_request`` with a live OTel span, ``function_setup``, then the limiter.""" + ``add_litellm_data_to_request`` with a live OTel span, ``function_setup``, then the limiter. + ``untagged_rate_limited_models`` are rejected before the limiter with a 429 that carries no + descriptor, the way the dynamic and batch limiters raise.""" from litellm.caching.caching import DualCache from litellm.proxy import proxy_server + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.utils import InternalUsageCache @@ -7115,6 +7120,11 @@ class TestPreCallWithFallbacksOnLocalRateLimit: skip_guardrails: bool = False, ) -> dict[str, object]: limiter_models.append(str(data["model"])) + if data["model"] in untagged_rate_limited_models: + raise ProxyRateLimitError( + detail=f"Priority rate limit exceeded for {data['model']}", + headers={"retry-after": "1"}, + ) await limiter.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=DualCache(), @@ -7146,37 +7156,49 @@ class TestPreCallWithFallbacksOnLocalRateLimit: @staticmethod def _otel_key( rpm_limit: int | None = None, + api_key: str = "hashed-key", model_rpm_limit: dict[str, int] | None = None, disable_fallbacks: bool | None = None, + team_model_rpm_limit: dict[str, int] | None = None, + organization_model_rpm_limit: dict[str, int] | None = None, + project_metadata: Mapping[str, Mapping[str, int]] | None = None, ) -> ProxyUserAPIKeyAuth: from opentelemetry.sdk.trace import TracerProvider - span = TracerProvider().get_tracer("test").start_span("proxy-request") + span: Final = TracerProvider().get_tracer("test").start_span("proxy-request") return ProxyUserAPIKeyAuth( - api_key="hashed-key", + api_key=api_key, parent_otel_span=span, rpm_limit=rpm_limit, metadata={ **({"model_rpm_limit": model_rpm_limit} if model_rpm_limit else {}), **({"disable_fallbacks": disable_fallbacks} if disable_fallbacks is not None else {}), }, + team_id="team-1" if team_model_rpm_limit else None, + team_metadata={"model_rpm_limit": team_model_rpm_limit} if team_model_rpm_limit else None, + org_id="org-1" if organization_model_rpm_limit else None, + organization_metadata=( + {"model_rpm_limit": organization_model_rpm_limit} if organization_model_rpm_limit else None + ), + project_id="project-1" if project_metadata else None, + project_metadata={k: dict(v) for k, v in project_metadata.items()} if project_metadata else None, ) @staticmethod def _chat_request() -> Request: return Request({"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": []}) - async def _pre_call( + async def _run( self, - data: dict[str, object], + processor: ProxyBaseLLMRequestProcessing, user_api_key_dict: ProxyUserAPIKeyAuth, rig: tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]], - ) -> tuple[ProxyBaseLLMRequestProcessing, tuple[dict[str, object], LiteLLMLoggingObj]]: + general_settings: Mapping[str, object] | None = None, + ) -> tuple[dict[str, object], LiteLLMLoggingObj]: proxy_logging_obj, router, proxy_config, _ = rig - processor = ProxyBaseLLMRequestProcessing(data=data) - result = await processor._pre_call_with_fallbacks( + return await processor._pre_call_with_fallbacks( request=self._chat_request(), - general_settings={}, + general_settings=dict(general_settings or {}), proxy_logging_obj=proxy_logging_obj, user_api_key_dict=user_api_key_dict, version=None, @@ -7190,7 +7212,16 @@ class TestPreCallWithFallbacksOnLocalRateLimit: route_type="acompletion", llm_router=router, ) - return processor, result + + async def _pre_call( + self, + data: dict[str, object], + user_api_key_dict: ProxyUserAPIKeyAuth, + rig: tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]], + general_settings: Mapping[str, object] | None = None, + ) -> tuple[ProxyBaseLLMRequestProcessing, tuple[dict[str, object], LiteLLMLoggingObj]]: + processor = ProxyBaseLLMRequestProcessing(data=data) + return processor, await self._run(processor, user_api_key_dict, rig, general_settings) @pytest.mark.asyncio async def test_v3_limiter_with_otel_span_falls_back_from_client_request(self, monkeypatch: pytest.MonkeyPatch): @@ -7265,6 +7296,232 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert processor.data["litellm_logging_obj"].model == primary_model assert processor.data["litellm_call_id"] + HARD_PER_MODEL_LIMITS: Final = MappingProxyType({"disable_fallbacks_on_per_model_rate_limits": True}) + + @pytest.mark.parametrize("cap_owner", ["key", "team"]) + @pytest.mark.asyncio + async def test_v3_limiter_per_model_cap_returns_429_when_per_model_limits_are_hard( + self, cap_owner: str, monkeypatch: pytest.MonkeyPatch + ): + """A key's own per-model cap and one inherited from its team both resolve into the key's + per-model descriptor (see ``get_key_model_rpm_limit``), so both must stop the fallback.""" + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model: Final = "gpt-4.1" + fallback_model: Final = "gpt-4.1-mini" + key: Final = ( + self._otel_key(model_rpm_limit={primary_model: 1}) + if cap_owner == "key" + else self._otel_key(team_model_rpm_limit={primary_model: 1}) + ) + rig: Final = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request: Final = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]} + + _, (first_data, _) = await self._pre_call(dict(request), key, rig, self.HARD_PER_MODEL_LIMITS) + processor: Final = ProxyBaseLLMRequestProcessing(data=dict(request)) + with pytest.raises(ProxyRateLimitError) as exc_info: + await self._run(processor, key, rig, self.HARD_PER_MODEL_LIMITS) + + assert first_data["model"] == primary_model + assert rig[3] == [primary_model, primary_model] + assert exc_info.value.status_code == 429 + assert exc_info.value.descriptor_key == "model_per_key" + assert exc_info.value.headers["retry-after"] + assert processor.data["model"] == primary_model + + @pytest.mark.asyncio + async def test_v3_limiter_team_cap_shared_by_two_keys_returns_429_when_per_model_limits_are_hard( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model: Final = "gpt-4.1" + fallback_model: Final = "gpt-4.1-mini" + first_key: Final = self._otel_key(team_model_rpm_limit={primary_model: 1}) + second_key: Final = self._otel_key(api_key="hashed-key-2", team_model_rpm_limit={primary_model: 1}) + rig: Final = self._v3_limiter_rig(monkeypatch, first_key, [{primary_model: [fallback_model]}]) + request: Final = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]} + + _, (first_data, _) = await self._pre_call(dict(request), first_key, rig, self.HARD_PER_MODEL_LIMITS) + processor: Final = ProxyBaseLLMRequestProcessing(data=dict(request)) + with pytest.raises(ProxyRateLimitError) as exc_info: + await self._run(processor, second_key, rig, self.HARD_PER_MODEL_LIMITS) + + assert first_data["model"] == primary_model + assert rig[3] == [primary_model, primary_model] + assert exc_info.value.status_code == 429 + assert exc_info.value.descriptor_key == "model_per_team" + assert exc_info.value.headers["retry-after"] + assert processor.data["model"] == primary_model + + @pytest.mark.parametrize("setting_value", [True, "true", "True", "1"]) + @pytest.mark.asyncio + async def test_per_model_limits_are_hard_for_true_and_a_true_string( + self, setting_value: object, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model: Final = "gpt-4.1" + fallback_model: Final = "gpt-4.1-mini" + key: Final = self._otel_key(model_rpm_limit={primary_model: 1}) + rig: Final = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request: Final = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]} + general_settings: Final = {"disable_fallbacks_on_per_model_rate_limits": setting_value} + + await self._pre_call(dict(request), key, rig, general_settings) + processor: Final = ProxyBaseLLMRequestProcessing(data=dict(request)) + with pytest.raises(ProxyRateLimitError) as exc_info: + await self._run(processor, key, rig, general_settings) + + assert exc_info.value.descriptor_key == "model_per_key" + assert rig[3] == [primary_model, primary_model] + assert processor.data["model"] == primary_model + + @pytest.mark.parametrize("setting_value", [False, "false", "", None, "not-a-bool"]) + @pytest.mark.asyncio + async def test_per_model_limits_stay_soft_for_every_other_setting_value( + self, setting_value: object, monkeypatch: pytest.MonkeyPatch + ): + primary_model: Final = "gpt-4.1" + fallback_model: Final = "gpt-4.1-mini" + key: Final = self._otel_key(model_rpm_limit={primary_model: 1}) + rig: Final = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request: Final = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]} + general_settings: Final = {"disable_fallbacks_on_per_model_rate_limits": setting_value} + + await self._pre_call(dict(request), key, rig, general_settings) + _, (data, _) = await self._pre_call(dict(request), key, rig, general_settings) + + assert data["model"] == fallback_model + assert rig[3] == [primary_model, primary_model, fallback_model] + + @pytest.mark.parametrize( + ("cap_owner", "expected_descriptor"), + [ + ("organization", "model_per_organization"), + ("project", "model_per_project"), + ("project_itpm", "model_per_project_itpm"), + ("project_otpm", "model_per_project_otpm"), + ], + ) + @pytest.mark.asyncio + async def test_v3_limiter_org_and_project_caps_return_429_when_per_model_limits_are_hard( + self, cap_owner: str, expected_descriptor: str, monkeypatch: pytest.MonkeyPatch + ): + """Every per-model descriptor the real limiter raises stops the fallback hunt, each under its + own key: an organization cap, a project cap, and the project's input and output token caps + (the token caps trip on the first request, since its own tokens already exceed a limit of 1).""" + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model: Final = "gpt-4.1" + fallback_model: Final = "gpt-4.1-mini" + key: Final = ( + self._otel_key(organization_model_rpm_limit={primary_model: 1}) + if cap_owner == "organization" + else self._otel_key( + project_metadata={ + { + "project": "model_rpm_limit", + "project_itpm": "model_itpm_limit", + "project_otpm": "model_otpm_limit", + }[cap_owner]: {primary_model: 1} + } + ) + ) + rig: Final = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request: Final = { + "model": primary_model, + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 8, + } + processors: Final = tuple(ProxyBaseLLMRequestProcessing(data=dict(request)) for _ in range(2)) + + async def run_both_requests() -> None: + for processor in processors: + await self._run(processor, key, rig, self.HARD_PER_MODEL_LIMITS) + + with pytest.raises(ProxyRateLimitError) as exc_info: + await run_both_requests() + + assert set(rig[3]) == {primary_model} + assert exc_info.value.status_code == 429 + assert exc_info.value.descriptor_key == expected_descriptor + assert exc_info.value.headers["retry-after"] + assert all(processor.data["model"] == primary_model for processor in processors) + + @pytest.mark.asyncio + async def test_v3_limiter_global_key_cap_still_tries_fallbacks_when_per_model_limits_are_hard( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model: Final = "gpt-4.1" + fallback_model: Final = "gpt-4.1-mini" + key: Final = self._otel_key(rpm_limit=1) + rig: Final = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request: Final = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]} + + await self._pre_call(dict(request), key, rig, self.HARD_PER_MODEL_LIMITS) + processor: Final = ProxyBaseLLMRequestProcessing(data=dict(request)) + with pytest.raises(ProxyRateLimitError) as exc_info: + await self._run(processor, key, rig, self.HARD_PER_MODEL_LIMITS) + + assert rig[3] == [primary_model, primary_model, fallback_model] + assert exc_info.value.descriptor_key == "api_key" + assert processor.data["model"] == primary_model + + @pytest.mark.asyncio + async def test_v3_limiter_fallback_per_model_cap_returns_429_when_per_model_limits_are_hard( + self, monkeypatch: pytest.MonkeyPatch + ): + """An untagged 429 on the primary still opens the fallback hunt, but a per-model cap on the + first fallback ends it with that fallback's 429 instead of moving on to the next model.""" + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model: Final = "gpt-4.1" + capped_fallback: Final = "gpt-4.1-mini" + last_fallback: Final = "gpt-4.1-nano" + key: Final = self._otel_key(model_rpm_limit={capped_fallback: 1}) + rig: Final = self._v3_limiter_rig( + monkeypatch, + key, + [{primary_model: [capped_fallback, last_fallback]}], + untagged_rate_limited_models=frozenset({primary_model}), + ) + request: Final = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]} + + _, (first_data, _) = await self._pre_call(dict(request), key, rig, self.HARD_PER_MODEL_LIMITS) + processor: Final = ProxyBaseLLMRequestProcessing(data=dict(request)) + with pytest.raises(ProxyRateLimitError) as exc_info: + await self._run(processor, key, rig, self.HARD_PER_MODEL_LIMITS) + + assert first_data["model"] == capped_fallback + assert rig[3] == [primary_model, capped_fallback, primary_model, capped_fallback] + assert exc_info.value.descriptor_key == "model_per_key" + assert processor.data["model"] == primary_model + + @pytest.mark.asyncio + async def test_v3_limiter_fallback_per_model_cap_moves_on_when_setting_is_off( + self, monkeypatch: pytest.MonkeyPatch + ): + primary_model: Final = "gpt-4.1" + capped_fallback: Final = "gpt-4.1-mini" + last_fallback: Final = "gpt-4.1-nano" + key: Final = self._otel_key(model_rpm_limit={capped_fallback: 1}) + rig: Final = self._v3_limiter_rig( + monkeypatch, + key, + [{primary_model: [capped_fallback, last_fallback]}], + untagged_rate_limited_models=frozenset({primary_model}), + ) + request: Final = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]} + + await self._pre_call(dict(request), key, rig) + _, (data, _) = await self._pre_call(dict(request), key, rig) + + assert data["model"] == last_fallback + assert rig[3] == [primary_model, capped_fallback, primary_model, capped_fallback, last_fallback] + @pytest.mark.asyncio async def test_fallback_lookup_uses_alias_resolved_model_group(self, monkeypatch: pytest.MonkeyPatch): primary_model = "gpt-4.1" diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 2ba736d8ae0..cf04b9ff91e 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -35,6 +35,7 @@ from litellm.litellm_core_utils.redact_messages import _get_turn_off_message_log from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.proxy._types import AddTeamCallback, ProxyException, TeamCallbackMetadata, UserAPIKeyAuth +from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES from litellm.proxy.litellm_pre_call_utils import ( KeyAndTeamLoggingSettings, LiteLLMProxyRequestSetup, @@ -1276,13 +1277,7 @@ async def test_key_metadata_enable_prompt_caching_promoted_to_request_root(key_v "service_callback", "logger_fn", "litellm_disabled_callbacks", - "_agentic_loop_depth", - "_agentic_loop_fingerprints", - "_code_interpreter_interception_active", - "_code_interpreter_interception_converted_stream", - "_code_interpreter_interception_sandbox_key", - "_headroom_interception_converted_stream", - "max_agentic_loops", + *AGENTIC_LOOP_KWARG_NAMES, ], ) async def test_add_litellm_data_to_request_strips_callback_control_fields( @@ -1313,13 +1308,18 @@ async def test_add_litellm_data_to_request_strips_callback_control_fields( "logger_fn": "module.func", "_agentic_loop_depth": 5, "_agentic_loop_fingerprints": ["forged"], + "_agentic_loop_api_surface": "responses", "_code_interpreter_interception_active": True, "_code_interpreter_interception_converted_stream": True, "_code_interpreter_interception_sandbox_key": "forged-key", + "_code_interpreter_interception_session_scoped": True, + "_websearch_interception_emit_native_blocks": True, + "_websearch_interception_converted_stream": True, + "_websearch_interception_stream_options": {"include_usage": True}, "_headroom_interception_converted_stream": True, "max_agentic_loops": 9999, } - sample_value = sample_values[control_field] + sample_value: Final = sample_values.get(control_field, "forged") updated = await add_litellm_data_to_request( data={ diff --git a/tests/unit/test_completion_timeout_resolution.py b/tests/unit/test_completion_timeout_resolution.py index 7eb79e90e60..77166d6e947 100644 --- a/tests/unit/test_completion_timeout_resolution.py +++ b/tests/unit/test_completion_timeout_resolution.py @@ -2,8 +2,11 @@ import os import sys +from typing import Final, Literal import httpx +import pytest +from openai import Timeout as SDKTimeout sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))) @@ -142,3 +145,35 @@ def test_httpx_timeout_preserved_for_openai(): ) assert out is t assert isinstance(out, httpx.Timeout) + + +@pytest.mark.parametrize("source", ("model", "timeout", "request_timeout")) +def test_sdk_timeout_preserves_components_for_httpx_providers( + source: Literal["model", "timeout", "request_timeout"], +) -> None: + timeout: Final = SDKTimeout(connect=2.0, read=None, write=5.0, pool=7.0) + resolved: Final = CompletionTimeout.resolve( + timeout if source == "model" else None, + {} if source == "model" else {source: timeout}, + "openai", + global_timeout=None, + supports_httpx_timeout=supports_httpx_timeout, + ) + assert isinstance(resolved, httpx.Timeout) + assert resolved.as_dict() == timeout.as_dict() + + +@pytest.mark.parametrize("read_timeout, expected", ((23.0, 23.0), (None, 600.0))) +def test_sdk_timeout_coerces_read_timeout_for_providers_without_httpx_support( + read_timeout: float | None, expected: float +) -> None: + assert ( + CompletionTimeout.resolve( + SDKTimeout(connect=2.0, read=read_timeout, write=5.0, pool=7.0), + {}, + "azure_ai", + global_timeout=None, + supports_httpx_timeout=supports_httpx_timeout, + ) + == expected + ) diff --git a/tests/unit/test_cost_map_guard.py b/tests/unit/test_cost_map_guard.py index 1595bd9864f..e127aa08d29 100644 --- a/tests/unit/test_cost_map_guard.py +++ b/tests/unit/test_cost_map_guard.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping import importlib.util import json import subprocess @@ -62,8 +63,19 @@ def _snapshot(cost_map: dict[str, object], backup: str | None = None, schema: st BASE: Final = _snapshot(BASE_MAP) -def _failures(head: object, changed_files: tuple[str, ...] = MAP_FILES, bot: bool = True) -> tuple[str, ...]: - return guard.guard_failures(BASE, head, changed_files, bot) +def _failures( + head: object, + changed_files: tuple[str, ...] = MAP_FILES, + bot: bool = True, + evaluation_output_price_sources: Mapping[str, str] = guard.EVALUATION_OUTPUT_PRICE_SOURCES, +) -> tuple[str, ...]: + return guard.guard_failures( + BASE, + head, + changed_files, + bot, + evaluation_output_price_sources=evaluation_output_price_sources, + ) def test_in_sync_files_pass_for_humans_and_bots() -> None: @@ -71,6 +83,52 @@ def test_in_sync_files_pass_for_humans_and_bots() -> None: assert _failures(BASE, bot=True) == () +EVALUATION_MODEL: Final = "azure_ai/Microsoft-Decision-1" +EVALUATION_OUTPUT_PRICE_FAILURE: Final = ( + f"{EVALUATION_MODEL}: 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)" +) + + +@pytest.mark.parametrize("bot", (True, False)) +def test_evaluation_rows_with_output_price_fail_for_humans_and_bots(bot: bool) -> None: + head: Final = _snapshot( + {**BASE_MAP, EVALUATION_MODEL: _entry(mode="evaluation", output_cost_per_token=4.2e-08)} + ) + assert _failures(head, bot=bot) == (EVALUATION_OUTPUT_PRICE_FAILURE,) + + +def test_evaluation_rows_with_zero_or_missing_output_price_pass() -> None: + zero_output: Final = _snapshot( + {**BASE_MAP, EVALUATION_MODEL: _entry(mode="evaluation", output_cost_per_token=0.0)} + ) + missing_output_entry: Final = { + key: value + for key, value in _entry(mode="evaluation").items() + if key != "output_cost_per_token" + } + missing_output: Final = _snapshot({**BASE_MAP, EVALUATION_MODEL: missing_output_entry}) + + assert _failures(zero_output) == () + assert _failures(missing_output) == () + + +def test_chat_rows_with_output_price_pass() -> None: + head: Final = _snapshot( + {**BASE_MAP, "openrouter/chat-with-output": _entry(output_cost_per_token=4.2e-08)} + ) + assert _failures(head) == () + + +def test_evaluation_output_price_source_allowlist_passes() -> None: + head: Final = _snapshot( + {**BASE_MAP, EVALUATION_MODEL: _entry(mode="evaluation", output_cost_per_token=4.2e-08)} + ) + source: Final = 'https://vendor.example/pricing "Output tokens are charged per token."' + + assert _failures(head, evaluation_output_price_sources={EVALUATION_MODEL: source}) == () + + def test_bot_may_add_and_reprice_models() -> None: head = _snapshot({**BASE_MAP, "openrouter/a": _entry(9e-06, supports_vision=True), "openrouter/c": _entry()}) assert _failures(head) == () diff --git a/tests/unit/test_exception_header_preservation.py b/tests/unit/test_exception_header_preservation.py index dd142d9d40b..7acacf05226 100644 --- a/tests/unit/test_exception_header_preservation.py +++ b/tests/unit/test_exception_header_preservation.py @@ -8,6 +8,8 @@ This is important for debugging and observability - headers like x-request-id, x-ms-region, rate limit headers, etc. should be available even when errors occur. """ +from typing import Final + import httpx import pytest @@ -17,6 +19,7 @@ from litellm.exceptions import ( ContextWindowExceededError, ImageFetchError, MidStreamFallbackError, + PaymentRequiredError, RateLimitError, ServiceUnavailableError, ) @@ -25,10 +28,11 @@ from litellm.exceptions import ( class TestExceptionHeaderPreservation: """Test that exception classes preserve headers from provider responses.""" - @pytest.fixture - def mock_response_with_headers(self) -> httpx.Response: + @pytest.fixture(params=("httpx", "httpx2")) + def mock_response_with_headers(self, request: pytest.FixtureRequest) -> httpx.Response: """Create a mock response with typical provider headers.""" - return httpx.Response( + transport: Final = pytest.importorskip(request.param) + return transport.Response( status_code=400, headers={ "x-request-id": "req-abc123", @@ -36,7 +40,7 @@ class TestExceptionHeaderPreservation: "x-ratelimit-remaining-requests": "99", "x-ratelimit-remaining-tokens": "9999", }, - request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions"), + request=transport.Request("POST", "https://api.openai.com/v1/chat/completions"), ) def test_bad_request_error_preserves_headers( @@ -50,11 +54,30 @@ class TestExceptionHeaderPreservation: response=mock_response_with_headers, ) - assert error.response is not None + assert error.response is mock_response_with_headers + assert error.request is mock_response_with_headers.request + assert error.request_id == mock_response_with_headers.headers["x-request-id"] assert error.response.headers.get("x-request-id") == "req-abc123" assert error.response.headers.get("x-ms-region") == "eastus" assert error.response.headers.get("x-ratelimit-remaining-requests") == "99" + def test_payment_required_error_preserves_headers( + self, mock_response_with_headers: httpx.Response + ): + """PaymentRequiredError should keep the provider response like its BadRequestError parent.""" + error = PaymentRequiredError( + message="Insufficient credit", + model="gpt-4", + llm_provider="openrouter", + response=mock_response_with_headers, + ) + + assert error.status_code == 402 + assert error.response is mock_response_with_headers + assert error.request is mock_response_with_headers.request + assert error.request_id == mock_response_with_headers.headers["x-request-id"] + assert error.response.headers.get("x-ms-region") == "eastus" + def test_content_policy_violation_error_preserves_headers( self, mock_response_with_headers: httpx.Response ): diff --git a/tests/unit/test_openai_embedding_encoding_format_default.py b/tests/unit/test_openai_embedding_encoding_format_default.py index 7a42eaf0f0a..9788eb30b1b 100644 --- a/tests/unit/test_openai_embedding_encoding_format_default.py +++ b/tests/unit/test_openai_embedding_encoding_format_default.py @@ -7,9 +7,19 @@ import respx import litellm +_SDK_EMBEDDING_MODELS: Final = ( + "openai/sdk-compat", + "mistral/sdk-compat", + "fireworks_ai/accounts/fireworks/models/sdk-compat", + "together_ai/sdk-compat", + "nvidia_nim/sdk-compat", +) -def _mock_openai_embedding_route(respx_mock: respx.MockRouter) -> respx.Route: - return respx_mock.post("https://api.openai.com/v1/embeddings").mock( + +def _mock_openai_embedding_route( + respx_mock: respx.MockRouter, api_base: str = "https://api.openai.com/v1" +) -> respx.Route: + return respx_mock.post(f"{api_base}/embeddings").mock( return_value=httpx.Response( 200, json={ @@ -27,13 +37,21 @@ def clear_default_encoding_format_env(monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) -def test_embedding_openai_omits_encoding_format_when_client_omits_it(respx_mock: respx.MockRouter) -> None: - mock_route: Final = _mock_openai_embedding_route(respx_mock) +@pytest.mark.parametrize("model", _SDK_EMBEDDING_MODELS) +def test_embedding_sdk_providers_omit_encoding_format_when_client_omits_it( + respx_mock: respx.MockRouter, model: str +) -> None: + provider, upstream_model = model.split("/", 1) + api_base: Final = f"https://{provider.replace('_', '-')}.example/v1" + mock_route: Final = _mock_openai_embedding_route(respx_mock, api_base) - response: Final = litellm.embedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + response: Final = litellm.embedding(model=model, input=["hello"], api_key="sk-test", api_base=api_base) request_body: Final = json.loads(mock_route.calls.last.request.read()) assert "encoding_format" not in request_body + assert request_body["model"] == upstream_model + assert request_body["input"] == ["hello"] + assert mock_route.calls.last.request.headers["authorization"] == "Bearer sk-test" assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] @@ -89,16 +107,22 @@ def test_embedding_openai_env_none_omits_encoding_format( @pytest.mark.asyncio -async def test_aembedding_openai_omits_encoding_format_when_client_omits_it( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +@pytest.mark.parametrize("model", _SDK_EMBEDDING_MODELS) +async def test_aembedding_sdk_providers_omit_encoding_format_when_client_omits_it( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, model: str ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - mock_route: Final = _mock_openai_embedding_route(respx_mock) + provider, upstream_model = model.split("/", 1) + api_base: Final = f"https://{provider.replace('_', '-')}.example/v1" + mock_route: Final = _mock_openai_embedding_route(respx_mock, api_base) - response: Final = await litellm.aembedding(model="openai/text-embedding-3-small", input=["hello"], api_key="sk-test") + response: Final = await litellm.aembedding(model=model, input=["hello"], api_key="sk-test", api_base=api_base) request_body: Final = json.loads(mock_route.calls.last.request.read()) assert "encoding_format" not in request_body + assert request_body["model"] == upstream_model + assert request_body["input"] == ["hello"] + assert mock_route.calls.last.request.headers["authorization"] == "Bearer sk-test" assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 034d501483c..467fe206089 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -53,7 +53,7 @@ from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.proxy.utils import is_valid_api_key from litellm.responses.main import aresponses, responses from litellm.types.caching import CachingSupportedCallTypes -from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY +from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY, STREAM_OPTIONS_STASH_KEYS from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams, LiteLLM_Params from litellm.types.utils import ( @@ -5291,6 +5291,7 @@ def test_function_setup_failure_log_line_shows_outer_not_doomed_ids(monkeypatch) WEBSEARCH_INTERNAL_CONTROL_FIELDS = ( "_websearch_interception_emit_native_blocks", "_websearch_interception_converted_stream", + "_websearch_interception_stream_options", ) @@ -5318,6 +5319,15 @@ def test_websearch_interception_control_fields_never_reach_the_provider(): assert set(WEBSEARCH_INTERNAL_CONTROL_FIELDS) <= set(all_litellm_params) +@pytest.mark.parametrize("stash_key", STREAM_OPTIONS_STASH_KEYS) +def test_converted_stream_options_stash_never_reaches_the_provider(stash_key: str): + non_default: Final = get_non_default_completion_params( + {"a_real_provider_specific_param": 1, stash_key: {"include_usage": True}} + ) + + assert non_default == {"a_real_provider_specific_param": 1} + + def test_get_litellm_params_keys_never_reach_the_provider(): """Bridges (chat <-> Responses, agentic loop follow-ups) forward litellm_params as `completion()` kwargs. Any key the param builder does not recognize is swept into diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 16c9b875636..d27dc733c78 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -227,9 +227,12 @@ AGENTIC_LOOP_STATE_NAMES: Final = ( "_code_interpreter_interception_sandbox_key", "_code_interpreter_interception_session_scoped", "_code_interpreter_interception_converted_stream", + "_code_interpreter_interception_stream_options", "_websearch_interception_emit_native_blocks", "_websearch_interception_converted_stream", + "_websearch_interception_stream_options", "_headroom_interception_converted_stream", + "_headroom_interception_stream_options", ) INTERNAL_STATE_NAMES: Final = ( diff --git a/tests/unit/types/test_router.py b/tests/unit/types/test_router.py index 89238e04a5a..6188bb966ee 100644 --- a/tests/unit/types/test_router.py +++ b/tests/unit/types/test_router.py @@ -1,6 +1,9 @@ import logging +from typing import Final +import httpx import pytest +from openai import Timeout as SDKTimeout from pydantic import ValidationError from litellm.types.router import ( @@ -26,6 +29,13 @@ from litellm.types.utils import ( ) +def test_sdk_timeout_is_normalized_for_provider_clients() -> None: + timeout: Final = SDKTimeout(connect=2.0, read=None, write=5.0, pool=7.0) + params: Final = GenericLiteLLMParams(timeout=timeout) + assert isinstance(params.timeout, httpx.Timeout) + assert params.timeout.as_dict() == timeout.as_dict() + + def test_model_info_declares_mirrored_pricing_fields(): """The pricing keys Deployment mirrors onto model_info must be declared fields, not extras that only survive because ModelInfo sets extra="allow".""" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index 55597f4ce82..a9cb68ae8dd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -22,15 +22,7 @@ vi.mock("./detailNavigation", () => ({ })); vi.mock("@/components/molecules/DecisionModelsBanner", () => ({ - default: ({ onAddModel }: { onAddModel?: () => void }) => ( -
- {onAddModel && ( - - )} -
- ), + default: () =>
, })); vi.mock("@/components/model_info_view", () => ({ default: ({ modelId }: { modelId: string }) =>
model:{modelId}
, @@ -107,21 +99,6 @@ describe("ModelsAndEndpointsPage", () => { expect(screen.queryByTestId("panel-all-models")).not.toBeInTheDocument(); }); - it("opens the Add Model tab from the decision models banner", async () => { - const user = userEvent.setup(); - renderPage(); - await user.click(screen.getByRole("button", { name: "Add a decision model" })); - expect(screen.getByTestId("panel-add")).toBeInTheDocument(); - expect(screen.queryByTestId("panel-all-models")).not.toBeInTheDocument(); - }); - - it("does not offer the banner's Add a decision model to a session that cannot add models", () => { - mockUseAuthorized.mockReturnValue(NON_ADMIN); - renderPage(); - expect(screen.getByRole("region", { name: "Decision models" })).toBeInTheDocument(); - expect(screen.queryByRole("button", { name: "Add a decision model" })).not.toBeInTheDocument(); - }); - it("does not show a view-only session the decision models banner, whose Playground it cannot open", () => { mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); renderPage(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index 1eae68b9d43..a54c4ee7292 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -178,9 +178,7 @@ export default function ModelsAndEndpointsPage() { - {!isViewOnly && ( - setActiveKey("add") : undefined} /> - )} + {!isViewOnly && } {modelId ? ( { "href", "https://docs.litellm.ai/docs/decisions", ); - expect(within(note).getByRole("link", { name: "test it in the System One playground" })).toHaveAttribute( + expect(within(note).getByRole("link", { name: "test it in the Decisions playground" })).toHaveAttribute( "href", "/ui/playground?tab=system-one", ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/JsonEditor.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/JsonEditor.tsx index 1effe55ef41..cc7bebb0bb8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/JsonEditor.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/JsonEditor.tsx @@ -1,11 +1,12 @@ import { Badge } from "@/components/ui/badge"; import { cn } from "@/lib/cva.config"; import { CircleAlert, CircleCheck, TriangleAlert } from "lucide-react"; -import { useId, useMemo, useRef } from "react"; +import { useTheme } from "next-themes"; +import { useId, useRef } from "react"; import { createElement, PrismLight as SyntaxHighlighter } from "react-syntax-highlighter"; import type { SyntaxHighlighterProps } from "react-syntax-highlighter"; import json from "react-syntax-highlighter/dist/esm/languages/prism/json"; -import { findRootBlocks, ROOT_BLOCK_STYLES, type RootBlock } from "./lib/rootBlocks"; +import { vs, vscDarkPlus } from "react-syntax-highlighter/dist/esm/styles/prism"; import type { SystemOnePayloadValidation } from "./lib/validatePayload"; SyntaxHighlighter.registerLanguage("json", json); @@ -14,16 +15,44 @@ const EDITOR_TEXT = "m-0 whitespace-pre-wrap wrap-anywhere py-3 font-mono text-x const GUTTER_WIDTH = "w-11"; type LineRendererProps = Parameters>[0]; const CONTENT_INSET = "pl-14 pr-3"; -const CODE_TAG_PROPS = { className: "language-json", style: { whiteSpace: "pre-wrap" } } as const; +type SyntaxTheme = NonNullable; +const THEME_BLOCKS = ['pre[class*="language-"]', 'code[class*="language-"]'] as const; -const TOKEN_COLORS = [ - "[&_.token.property]:text-sky-700 dark:[&_.token.property]:text-sky-300", - "[&_.token.string]:text-emerald-700 dark:[&_.token.string]:text-emerald-300", - "[&_.token.number]:text-amber-700 dark:[&_.token.number]:text-amber-300", - "[&_.token.boolean]:text-violet-700 dark:[&_.token.boolean]:text-violet-300", - "[&_.token.null]:text-violet-700 dark:[&_.token.null]:text-violet-300", - "[&_.token.punctuation]:text-muted-foreground [&_.token.operator]:text-muted-foreground", -].join(" "); +function withoutBlockBackground(theme: SyntaxTheme): SyntaxTheme { + const stripped = THEME_BLOCKS.map((block) => { + const { background: _background, backgroundColor: _backgroundColor, ...rest } = theme[block] ?? {}; + return [block, rest] as const; + }); + return { ...theme, ...Object.fromEntries(stripped) }; +} + +const VS_CODE_DARK_PLUS = withoutBlockBackground(vscDarkPlus); +const VS_CODE_LIGHT_PLUS = { + ...withoutBlockBackground(vs), + property: { color: "#0451a5" }, + string: { color: "#a31515" }, + number: { color: "#098658" }, + boolean: { color: "#0000ff" }, + punctuation: { color: "#000000" }, + operator: { color: "#000000" }, +} as const; +const INHERIT_TEXT = { + fontFamily: "inherit", + fontSize: "inherit", + lineHeight: "inherit", + whiteSpace: "pre-wrap", + overflowWrap: "anywhere", + textShadow: "none", +} as const; +const CODE_TAG_PROPS = { className: "language-json", style: INHERIT_TEXT } as const; +const TRANSPARENT_PRE = { + ...INHERIT_TEXT, + background: "transparent", + border: 0, + margin: 0, + padding: 0, + overflow: "visible", +} as const; interface JsonEditorProps { value: string; @@ -31,7 +60,7 @@ interface JsonEditorProps { validation: SystemOnePayloadValidation; } -function ValidationStatus({ validation }: { validation: SystemOnePayloadValidation }) { +export function ValidationStatus({ validation }: { validation: SystemOnePayloadValidation }) { const errorCount = validation.issues.filter((issue) => issue.severity === "error").length; if (errorCount > 0) { return ( @@ -43,7 +72,7 @@ function ValidationStatus({ validation }: { validation: SystemOnePayloadValidati return Valid payload; } -function IssueList({ id, validation }: { id: string; validation: SystemOnePayloadValidation }) { +export function IssueList({ id, validation }: { id: string; validation: SystemOnePayloadValidation }) { if (validation.issues.length === 0) { return (

@@ -69,39 +98,24 @@ function IssueList({ id, validation }: { id: string; validation: SystemOnePayloa ); } -function renderLines(rootBlocks: readonly RootBlock[]) { - return function LineRows({ rows, stylesheet, useInlineStyles }: LineRendererProps) { - return rows.map((row, line) => { - const lineElement = { node: row, stylesheet, useInlineStyles, key: line }; - const block = rootBlocks.find(({ startLine, endLine }) => line >= startLine && line <= endLine); - return ( -

- - {line + 1} - - - {createElement(lineElement)} - -
- ); - }); - }; +function LineRows({ rows, stylesheet, useInlineStyles }: LineRendererProps) { + return rows.map((row, line) => { + const lineElement = { node: row, stylesheet, useInlineStyles, key: line }; + return ( +
+ + {line + 1} + + {createElement(lineElement)} +
+ ); + }); } export default function JsonEditor({ value, onChange, validation }: JsonEditorProps) { const issuesId = useId(); const highlightRef = useRef(null); - const renderer = useMemo(() => renderLines(findRootBlocks(value)), [value]); + const syntaxTheme = useTheme().resolvedTheme === "dark" ? VS_CODE_DARK_PLUS : VS_CODE_LIGHT_PLUS; const lineCount = value.split("\n").length; const hasErrors = !validation.isValid; @@ -121,20 +135,22 @@ export default function JsonEditor({ value, onChange, validation }: JsonEditorPr {lineCount} {lineCount === 1 ? "line" : "lines"} -
-